{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NNX\n",
    "\n",
    "Welcome to NNX!\n",
    "\n",
    "NNX is an open source Python library for **N**eural **N**etwork in JA**X**. Its main feature is, much like Pytorch, allowing Python object semantics and reference sharing, which brings simplicty and familiarity, and easily crossing over into the functional world with through a set of simple APIs.\n",
    "\n",
    "This tutorial demonstrates how to construct a simple convolutional neural network (CNN) using NNX and train the network for image classification on the MNIST dataset."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -q nnx"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the MNIST dataset\n",
    "We will use the `datasets` library to load MNIST and convert it to NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cris/nnx/.venv/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Found cached dataset mnist (/home/cris/.cache/huggingface/datasets/mnist/mnist/1.0.0/9d494b7f466d6931c64fb39d58bb1249a4d85c9eb9865d9bc20960b999e2a332)\n",
      "100%|██████████| 2/2 [00:00<00:00, 499.95it/s]\n"
     ]
    }
   ],
   "source": [
    "import datasets\n",
    "import numpy as np\n",
    "\n",
    "dataset = datasets.load_dataset(\"mnist\")\n",
    "X_train = np.array(np.stack(dataset[\"train\"][\"image\"]), dtype=np.uint8)[\n",
    "    ..., None\n",
    "]\n",
    "y_train = np.array(dataset[\"train\"][\"label\"], dtype=np.uint8)\n",
    "X_test = np.array(np.stack(dataset[\"test\"][\"image\"]), dtype=np.uint8)[..., None]\n",
    "y_test = np.array(dataset[\"test\"][\"label\"], dtype=np.uint8)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets visualize a few examples from the dataset using matplotlib:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAeQAAAH4CAYAAACbup4ZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAA58klEQVR4nO3df3zNdf/H8dcxs83P+f0rTWtcfi3UjESGvi3R1apFP7TQD3XxbYmkK2zlilQi+VkpikTzI+RSucxVXNpIiJIRlVVsFiM/Ztvn+0dfuzrn9WFn29n2PmeP++3mj/fT+3zO23q3l4/z2vvjsCzLEgAAUK4qlfcCAAAABRkAACNQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADAABRkAAANQkAEAMAAFWUQOHTokDodDXn75ZY9dc+PGjeJwOGTjxo0euyYqBvYjTMJ+LDteW5Dnz58vDodDtm3bVt5LKTXr16+Xnj17Sr169SQ4OFgiIyPl3XffLe9lwYav78fvvvtORowYIV27dpXAwEBxOBxy6NCh8l4WLsLX92NiYqI4HA71KzAwsLyXViKVy3sBsLdq1SqJiYmRa6+9tmDzLV26VOLi4iQzM1NGjBhR3ktEBbJlyxaZPn26tGnTRlq3bi07duwo7yUBMnv2bKlevXrB2M/PrxxXU3IUZEPNmDFDGjduLBs2bJCAgAARERk6dKi0atVK5s+fT0FGmfrrX/8qx48flxo1asjLL79MQYYRYmNjpV69euW9DI/x2n+ydkdOTo6MHz9errnmGqlVq5ZUq1ZNunfvLsnJyRd9zdSpUyUkJESCgoKkR48esnv3bjVn7969EhsbK3Xq1JHAwECJiIiQVatWFbqe06dPy969eyUzM7PQudnZ2VK7du2CYiwiUrlyZalXr54EBQUV+nqYx5v3Y506daRGjRqFzoP38Ob9eIFlWZKdnS2+8tBCny7I2dnZ8uabb0pUVJRMnjxZEhMTJSMjQ6Kjo23/hv/OO+/I9OnTZdiwYfL000/L7t27pVevXnLkyJGCOXv27JEuXbrIt99+K2PGjJEpU6ZItWrVJCYmRlasWHHJ9aSmpkrr1q1lxowZha49KipK9uzZI+PGjZP9+/fLgQMHZMKECbJt2zYZPXp0kb8WKH/evB/he3xhP4aGhkqtWrWkRo0aMnDgQKe1eCXLS7399tuWiFhbt2696Jzc3Fzr3LlzTtlvv/1mNWzY0BoyZEhBdvDgQUtErKCgIOvw4cMFeUpKiiUi1ogRIwqy3r17W+Hh4dbZs2cLsvz8fKtr165WixYtCrLk5GRLRKzk5GSVJSQkFPrnO3XqlNW/f3/L4XBYImKJiFW1alVr5cqVhb4WZc/X9+OfvfTSS5aIWAcPHizS61B2fH0/Tps2zRo+fLi1aNEiKykpyYqPj7cqV65stWjRwjpx4kShrzeVT98h+/n5SZUqVUREJD8/X7KysiQ3N1ciIiJk+/btan5MTIw0bdq0YBwZGSmdO3eWtWvXiohIVlaWbNiwQfr37y8nT56UzMxMyczMlGPHjkl0dLSkpaVJenr6RdcTFRUllmVJYmJioWsPCAiQli1bSmxsrCxevFgWLlwoERERMnDgQPniiy+K+JWACbx5P8L3ePN+jI+Pl9dee03uueceueOOO2TatGmyYMECSUtLk1mzZhXxK2EOny7IIiILFiyQq666SgIDA6Vu3bpSv359+eijj+TEiRNqbosWLVTWsmXLgh/v2L9/v1iWJePGjZP69es7/UpISBARkaNHj3pk3cOHD5fVq1fL+++/L3fddZfce++9sn79emncuLHEx8d75D1Q9rx1P8I3+dJ+vOeee6RRo0ayfv36UnuP0ubTXdYLFy6UQYMGSUxMjDz55JPSoEED8fPzk0mTJsmBAweKfL38/HwRERk1apRER0fbzgkLCyvRmkX+aLaYN2+ejB49WipV+u/fmfz9/aVPnz4yY8YMycnJKfjbLbyDt+5H+CZf3I/NmjWTrKysUn2P0uTTBTkpKUlCQ0Nl+fLl4nA4CvILf1tzlZaWprJ9+/ZJ8+bNReSPBgKRPwrjDTfc4PkF/79jx45Jbm6u5OXlqd87f/685Ofn2/4ezOat+xG+ydf2o2VZcujQIenYsWOZv7en+PQ/WV/4IXHrTy3xKSkpsmXLFtv5K1eudPqMIzU1VVJSUqRPnz4iItKgQQOJioqSuXPnyi+//KJen5GRccn1uNvW36BBAwkODpYVK1ZITk5OQX7q1ClZvXq1tGrVih998kLeuh/hm7x5P9pda/bs2ZKRkSE33XRToa83ldffIb/11luybt06lcfHx0u/fv1k+fLlctttt0nfvn3l4MGDMmfOHGnTpo2cOnVKvSYsLEy6desmjz76qJw7d06mTZsmdevWdfoxo5kzZ0q3bt0kPDxcHnroIQkNDZUjR47Ili1b5PDhw7Jz586LrjU1NVV69uwpCQkJl2xc8PPzk1GjRsnYsWOlS5cuEhcXJ3l5eTJv3jw5fPiwLFy4sGhfJJQZX9yPIiInTpyQ1157TURENm/eLCJ/HF4THBwswcHBMnz4cHe+PChjvrofQ0JCZMCAARIeHi6BgYGyadMmef/996VDhw4ydOhQ979Apimv9u6SutDWf7FfP/30k5Wfn29NnDjRCgkJsQICAqyOHTtaa9asse6//34rJCSk4FoX2vpfeukla8qUKVazZs2sgIAAq3v37tbOnTvVex84cMCKi4uzGjVqZPn7+1tNmza1+vXrZyUlJRXM8cSPmSxatMiKjIy0goODraCgIKtz585O7wFz+Pp+vLAmu19/XjvM4Ov78cEHH7TatGlj1ahRw/L397fCwsKsp556ysrOzi7Jl63cOSzLR444AQDAi/n0Z8gAAHgLCjIAAAagIAMAYAAKMgAABqAgAwBgAAoyAAAGoCADAGAAt0/q+vNZp4Crsv5xdvYjLoX9CJO4ux+5QwYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxQubwXAKBoPv/8c5Vdd911Ktu0aZPKbrvtNpUdO3bMMwsDUCLcIQMAYAAKMgAABqAgAwBgAAoyAAAG8NmmrgYNGqjsySefdBqPGjXKrWtNnz5dZc8884zKTp065ebqgOKzLMutrFu3biqbMWOGyu6++27PLAwoR127dlVZYmKiyiIiIlQWGRmpsv3793tkXUXBHTIAAAagIAMAYAAKMgAABqAgAwBgAJ9t6oqPj1fZyJEjncZ2jTB2/vd//1dloaGhKhswYIDKTp8+7dZ7AGWhSZMm5b0EoEDNmjVV5nrqXHh4uJozePBglTVv3lxlVapUcWsddk3ANHUBAFBBUZABADAABRkAAANQkAEAMIBPNHXZnTT08MMPF/q6pKQklc2dO1dlderUUdmcOXNUtmTJEpXFxsaq7Ny5c4WuDQDKUlhYmNO4bdu2xb5Ww4YNVdanTx+V2Z2Q1ahRo0Kv73A4VGb3ffXDDz9U2eLFi1W2d+/eQt+zLHCHDACAASjIAAAYgIIMAIABKMgAABjAJ5q60tPTVVa3bl2Vbd261Wlsd7KWu6d3tWnTRmUJCQkq+9vf/qayqVOnuvUeAFBSwcHBKvvoo49UdtVVVzmNq1atWuz3tGu6cvd769GjR53GKSkpas7mzZtV9sEHH6js0KFDbr2nKbhDBgDAABRkAAAMQEEGAMAAFGQAAAzgE01dV155ZbFe526TgZ3Jkyer7K677lJZpUr8nQdA+Zk3b57Krr32WpW58/1w0aJFKnP35MHly5er7MSJEyr74YcfnMZ2Tbu+imoBAIABKMgAABiAggwAgAEoyAAAGMAnmrree+89ld10000qu/rqq53GHTp0UHN27Njh1nuePXtWZd98843KittwBgCeYNc4ZXeS1sqVK53Gt99+e2ktCRfBHTIAAAagIAMAYAAKMgAABvCJz5DtfjD9mWeeUVlycrLT+LPPPlNzbr311kJfJyJSs2ZNlf31r39V2caNG1UGlAW7zwnh25o0aaKyqKgoldkdArJhwwansWvPjYjI5ZdfrrJ9+/apzPWJTReTnZ2tspycHLde64u4QwYAwAAUZAAADEBBBgDAABRkAAAM4BNNXXb279+vsri4OKfxqlWr1JwVK1aozK7Rq3nz5irz8/NT2bp16y61TOCSBg0apLLw8HC3XmvXuFORnpzj6xo3bqwyu+83ISEhbl3v1VdfLfGaLrBrKLTbj6mpqSr76KOPnMazZs1Sc7KyskqwOnNxhwwAgAEoyAAAGICCDACAASjIAAAYwGebuuy4nrg1ZMgQNWfJkiUqc30Kit21gJKKiYlR2ezZs1VWpUoVt65ndxLd8OHDi7wulL9WrVqpbOnSpSpr27atW9f75ZdfVHbkyBGn8eLFi91cnTZ48GCV2TV1tWnTRmWRkZFO4379+qk5H3zwgcpee+01lXnbqV/cIQMAYAAKMgAABqAgAwBgAAoyAAAGqFBNXa6WLVumsquuukpldqfJ2DXg2PHVE2XgeZ07d1ZZYGCgyuyaY+zYPXYP5rP7HrR582aVVa1aVWXHjx9X2QMPPKCyLVu2qMy1qaskXnrpJbfm2TWrRUdHO42feuopt64fFhamskcffdStdZiCO2QAAAxAQQYAwAAUZAAADEBBBgDAABW6qSs/P19lu3fvVllSUpLKBg4c6NZ7pKWlFX1hqJDsmrXs9qgdu8eNwnfk5eWp7P3331eZXQPX2bNnS2VNnrB3795Cs3/+859qzptvvqmy+++/X2Vz585V2Y4dO4qwwrLFHTIAAAagIAMAYAAKMgAABqAgAwBggArd1OWuatWqFfu1M2fOVNnjjz/uNOZRjiipqVOnlvcS4CG7du1Smd3pXT/++GNZLKfc7du3T2WTJk1S2Zo1a1Tm+r1WRGTQoEGeWFap4A4ZAAADUJABADAABRkAAANQkAEAMABNXW7o3bu3W/Pmz5+vsjvvvFNlQ4YMcRrT1AXgUipKA5e7srOz3ZpXs2bNUl6JZ3GHDACAASjIAAAYgIIMAIAB+AzZReXK+kvicDhUlpOTo7Lp06erzO6H2hMTE53Gzz//vJpj9xQU+JY6deo4jaOjo916XXp6uso+//xzj6wJ8AadOnUq7yWUCu6QAQAwAAUZAAADUJABADAABRkAAAPQ1OXilltuUVmNGjVUtmfPHpXt2LFDZQcPHlTZTTfd5DSeMmWKmtO3b99LLRM+ICsry2n88ccfqzkdOnRQWdOmTVXWvXt3ldntUXgnu6c9NWjQQGXr168vi+WUu3bt2rk1b+3ataW8Es/iDhkAAANQkAEAMAAFGQAAA1CQAQAwAE1dxeRu88SJEydUtmTJEqfxK6+8ouaEhYWpbP/+/W6uDhXNrbfeqrI5c+aUw0pQGqpXr66yTz75RGW9evVS2caNG0tjSWXmjjvuUNkDDzygMrsnQNl9jUzGHTIAAAagIAMAYAAKMgAABqAgAwBgAJq6XLRs2dKteZ5ssAoICFBZeHh4qb4nzLNr1y6VnT9/XmX+/v4qu/HGG1W2evVqlb333nsqy83NdRp/8MEHl1wnyt7OnTtVtnz5cpW9//77KrNr9Prmm288s7BSMGDAAKfxm2++qebYPf7Wrnntxx9/9Ni6ygJ3yAAAGICCDACAASjIAAAYgIIMAIABaOpysW/fPrfm2T3ububMmSqzO2GnX79+TuPMzEw1x9tP10HR2TXk1K9fX2UTJ05UWdWqVVVm9wjPm2++WWV5eXlO4zp16qg5c+fOVRnKzu+//66y4cOHq8zuEZ5ff/21yuwapV599VWnsacbv3r37q2yNm3aqOzFF190GlepUkXNsft/5d577y3B6szAHTIAAAagIAMAYAAKMgAABqAgAwBgAIdlWZZbEx2O0l6LEYKCglT21Vdfqczu8YgLFy5UWdOmTVXm2txgdzKS62k1pnNzG3lMRdmPdmJiYlT2+OOPq8z1MZ8i7v13stvvKSkpbq3NFBV1PzZp0kRln3/+ucquuOIKlbk2jp05c6bY67D7etSqVUtlfn5+Kjt69KjTePz48WrOu+++q7KzZ88WZYllyt39yB0yAAAGoCADAGAACjIAAAagIAMAYACautxgd9pWQkKCymJjY1VWs2ZNlS1btsxp/Mwzz6g5GRkZRVliuauoTTQwE/vxv9q1a6ey+Ph4lbk+8rVTp07Ffk+7r0dSUpLKdu/erbJ58+Y5jdPT04u9DlPQ1AUAgBehIAMAYAAKMgAABuAzZHgEn9nBJOxHmITPkAEA8CIUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADOD24xcBAEDp4Q4ZAAADUJABADAABRkAAANQkAEAMAAFWUQOHTokDodDXn75ZY9dc+PGjeJwOGTjxo0euyYqBvYjTMJ+LDteW5Dnz58vDodDtm3bVt5LKRWJiYnicDjUr8DAwPJeGmz4+n5csWKFREdHS5MmTSQgIEAuu+wyiY2Nld27d5f30mDD1/ejiMj69eulZ8+eUq9ePQkODpbIyEh59913y3tZJVK5vBeAS5s9e7ZUr169YOzn51eOq0FF9fXXX0vt2rUlPj5e6tWrJ7/++qu89dZbEhkZKVu2bJH27duX9xJRgaxatUpiYmLk2muvLbh5Wbp0qcTFxUlmZqaMGDGivJdYLBRkw8XGxkq9evXKexmo4MaPH6+yBx98UC677DKZPXu2zJkzpxxWhYpqxowZ0rhxY9mwYYMEBASIiMjQoUOlVatWMn/+fK8tyF77T9buyMnJkfHjx8s111wjtWrVkmrVqkn37t0lOTn5oq+ZOnWqhISESFBQkPTo0cP2n+T27t0rsbGxUqdOHQkMDJSIiAhZtWpVoes5ffq07N27VzIzM93+M1iWJdnZ2cL5Ld7PF/bjnzVo0ECqVq0qx48fL9brUb68eT9mZ2dL7dq1C4qxiEjlypWlXr16EhQUVOjrTeXTBTk7O1vefPNNiYqKksmTJ0tiYqJkZGRIdHS07NixQ81/5513ZPr06TJs2DB5+umnZffu3dKrVy85cuRIwZw9e/ZIly5d5Ntvv5UxY8bIlClTpFq1ahITEyMrVqy45HpSU1OldevWMmPGDLf/DKGhoVKrVi2pUaOGDBw40Gkt8C6+sB+PHz8uGRkZ8vXXX8uDDz4o2dnZ0rt3b7dfD3N4836MioqSPXv2yLhx42T//v1y4MABmTBhgmzbtk1Gjx5d5K+FMSwv9fbbb1siYm3duvWic3Jzc61z5845Zb/99pvVsGFDa8iQIQXZwYMHLRGxgoKCrMOHDxfkKSkplohYI0aMKMh69+5thYeHW2fPni3I8vPzra5du1otWrQoyJKTky0RsZKTk1WWkJBQ6J9v2rRp1vDhw61FixZZSUlJVnx8vFW5cmWrRYsW1okTJwp9PcqWr+/HC/7yl79YImKJiFW9enVr7NixVl5entuvR9nw9f146tQpq3///pbD4SjYj1WrVrVWrlxZ6GtN5tN3yH5+flKlShUREcnPz5esrCzJzc2ViIgI2b59u5ofExMjTZs2LRhHRkZK586dZe3atSIikpWVJRs2bJD+/fvLyZMnJTMzUzIzM+XYsWMSHR0taWlpkp6eftH1REVFiWVZkpiYWOja4+Pj5bXXXpN77rlH7rjjDpk2bZosWLBA0tLSZNasWUX8SsAE3rwfL3j77bdl3bp1MmvWLGndurWcOXNG8vLy3H49zOHN+zEgIEBatmwpsbGxsnjxYlm4cKFERETIwIED5YsvvijiV8Ig5fwXgmJz52+AlmVZ8+fPt8LDwy1/f/+Cv0mJiHXFFVcUzLnwN8Dx48er1993331WQECAZVn//RvhpX5t377dsiz7vwF6QqNGjazevXt79JoouYq4H7OysqyGDRtaI0eO9Ng14Rm+vh+HDh1qtW/f3ulfZ3JycqwWLVpYkZGRxbqmCXy6y3rhwoUyaNAgiYmJkSeffFIaNGggfn5+MmnSJDlw4ECRr5efny8iIqNGjZLo6GjbOWFhYSVac2GaNWsmWVlZpfoeKB2+th9r164tvXr1kkWLFnn00AiUDW/djzk5OTJv3jwZPXq0VKr033/k9ff3lz59+siMGTMkJyen4O7fm/h0QU5KSpLQ0FBZvny5OByOgjwhIcF2flpamsr27dsnzZs3F5E/GqxE/vgPf8MNN3h+wYWwLEsOHTokHTt2LPP3Rsn52n4UETlz5oycOHGiXN4bJeOt+/HYsWOSm5tr+1HJ+fPnJT8/32s/RvH5z5BFxOlHhlJSUmTLli2281euXOn0GUdqaqqkpKRInz59ROSPH/OIioqSuXPnyi+//KJen5GRccn1FKWt3+5as2fPloyMDLnpppsKfT3M48378ejRoyo7dOiQ/Otf/5KIiIhCXw/zeOt+bNCggQQHB8uKFSskJyenID916pSsXr1aWrVq5bU/+uT1d8hvvfWWrFu3TuXx8fHSr18/Wb58udx2223St29fOXjwoMyZM0fatGkjp06dUq8JCwuTbt26yaOPPirnzp2TadOmSd26dZ3a6GfOnCndunWT8PBweeihhyQ0NFSOHDkiW7ZskcOHD8vOnTsvutbU1FTp2bOnJCQkFNq4EBISIgMGDJDw8HAJDAyUTZs2yfvvvy8dOnSQoUOHuv8FQpny1f0YHh4uvXv3lg4dOkjt2rUlLS1N5s2bJ+fPn5cXXnjB/S8QypQv7kc/Pz8ZNWqUjB07Vrp06SJxcXGSl5cn8+bNk8OHD8vChQuL9kUySfl+hF18F5oWLvbrp59+svLz862JEydaISEhVkBAgNWxY0drzZo11v3332+FhIQUXOtC08JLL71kTZkyxWrWrJkVEBBgde/e3dq5c6d67wMHDlhxcXFWo0aNLH9/f6tp06ZWv379rKSkpII5JW3rf/DBB602bdpYNWrUsPz9/a2wsDDrqaeesrKzs0vyZUMp8fX9mJCQYEVERFi1a9e2KleubDVp0sS66667rF27dpXky4ZS4uv70bIsa9GiRVZkZKQVHBxsBQUFWZ07d3Z6D2/ksCyOgAIAoLz59GfIAAB4CwoyAAAGoCADAGAACjIAAAagIAMAYAAKMgAABqAgAwBgALdP6vrzWaeAq7L+cXb2Iy6F/QiTuLsfuUMGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMULm8FwBUBA888IDK3njjDY9d3+FwqOzw4cMqmzhxospmz57tsXUAKD7ukAEAMAAFGQAAA1CQAQAwAAUZAAAD0NQFlIE+ffqozLIsj13f7lpNmjRR2fTp01XWoUMHp/HQoUM9ti6gqPz8/JzGHTt2VHOef/55ldnNu+eee1S2fv36EqyudHGHDACAASjIAAAYgIIMAIABKMgAABiApq4KrG7duk7jVq1aufW6zZs3l8ZyfNqXX36psmuuucZpfPnll5f6OipV0n8Hdz1FbM+ePWqOXTMYUFKuDYUiIgsWLHAaX3XVVcW+frNmzYr92vLAHTIAAAagIAMAYAAKMgAABqAgAwBgAJq6fFCNGjVUNmTIEJWNHDnSady0aVO3ru96kg4KN2nSJJV98MEHTuNGjRp59D1feeUVlbk2konoRze6NvsBnjB48GCVPfPMMyq78sorncb/+c9/1JzTp0+r7OzZsypbuHBhUZZY7rhDBgDAABRkAAAMQEEGAMAAFGQAAAxAU5cXadu2rcoee+wxlUVHR6usuCfW/PDDD8V6HQq3f//+S46Lwu6/77lz54p9PcBddk2e48ePV9kTTzyhssqVdQl6+umnncavv/66mrNx40aVNWzYUGVVqlRR2fnz51VmCu6QAQAwAAUZAAADUJABADAABRkAAAPQ1GUAu8ceDho0SGUPPfSQyoKDg1U2b948lS1btkxlgYGBTuPevXurOX//+99VhrJj15QyYcIEldmdguTuiVu///6703j16tVurg4QeeSRR1Rm19S1detWldmdIPjdd985jZcvX67mhIeHq+yNN95QmeveNh13yAAAGICCDACAASjIAAAYgIIMAIABaOoqZS1btlTZmDFjnMZ33nmnmlO1alWVLVq0SGV2DQ8rV64swgr/68MPPyzW61B6VqxYobKbbrrJo+/h+gi8bdu2efT68B133323yl544QWV/etf/1LZAw88oDK7kwATExOdxv369VNzsrOzVbZkyRKVeRvukAEAMAAFGQAAA1CQAQAwQIX+DNnu6SDPP/+8ys6cOaOyhIQEld1zzz0qc/18TkR/Pvzuu++qOXafDa9fv15l8A59+/ZVmevTu5577jk1x9/fv9jvmZGRobI5c+aobNasWcV+D1QsdocY2X2e+/jjj6vM7vPiW265RWWuPTY5OTlqTmxsrMrsPrf2NtwhAwBgAAoyAAAGoCADAGAACjIAAAaoUE1drk+/WbdunZrTvn17lVmWpbJbb71VZfXr11eZ3SELzz77rNOYZi3v1bRpU5Xdd999Khs3bpzKXJ+25a5Dhw6pbMCAASr78ccfVXb06NFivScgIlKnTh2VJSUlqWz37t0qi46OVtmUKVMKfU+7xthPP/200Nd5I+6QAQAwAAUZAAADUJABADAABRkAAANUqKau0aNHO42vuuqqYl/L4XCo7Pbbb1fZP//5z2K/B8zXsWNHldmd9lZcmzZtUtmLL76oMp7QhLJw8uRJlT388MMqs2v+snuqnZ3Bgwc7jRcvXuzm6rwfd8gAABiAggwAgAEoyAAAGICCDACAAXy2qatPnz4qe+KJJ5zGx48fV3Nq166tsm+++UZld911l8r27NlThBXCF9g9eu73339XWbVq1Yp1/W7duqksLCxMZXZ7Lz09XWWvvvpqofPsHtsIiNh/z3Q9AVFEZODAgW5db9CgQSqrSE1crrhDBgDAABRkAAAMQEEGAMAAFGQAAAzgs01dMTExKqtUyfnvH8uXL1dz3n77bZV99dVXKjtz5kzxFwef8dlnn6nM7sS2ESNGqMz1pLgmTZq49Z6NGjVyK7MTFxensu3btzuN9+3bp+aMHDlSZb/++qtb7wmIiPzjH/9QWUVu4LLDHTIAAAagIAMAYAAKMgAABqAgAwBgAJ9t6tq/f7/KXB+Z2L17dzXnoYceKrU1oWJYv369W1mVKlWcxkOHDlVzJk+erLKAgIASrE67+uqrLzkWEWnXrp3KXnvtNZW99dZbKsvPzy/B6lBemjdvrrIHH3yw2Nf7+uuvVZaTk1Ps6/ki7pABADAABRkAAANQkAEAMAAFGQAAA/hsU9fu3btVlpeX5zRu0aKFmhMREaGybdu2eW5hwP9zbWixa5LaunWryuweLWrH7vSu4jbl2DV1zZ07V2X169dX2aRJk4r1nihbrk2GL7/8sppj9z3z448/Vtn111+vsmuvvVZlS5cuLcoSfR53yAAAGICCDACAASjIAAAYwGFZluXWRJdDNbzRCy+84DR+8skn1ZzffvtNZe3bt1dZenq65xbmA9zcRh5jyn60+zy3R48eKnv99ddV9v3335fKmi7w8/NTWdWqVVV26623Oo2HDx+u5nTq1Mmt93Tt0xARueWWW1Rm97mjJ1XU/VgSTz31lNPY9fuliMj8+fNVZnegzXPPPacyu/9X7L63+iJ39yN3yAAAGICCDACAASjIAAAYgIIMAIABfPZgEDsvvvii0/jOO+9Uc6644gqVVa9evdTWBO9m1zhl1yyYkpKistJu6rJrsDp58qTKFi5c6DS2e1Lahg0bVGb31Cm7r0flyhXq24xXuOGGG1Q2YcIEp3FaWpqaM2bMGJXZPbHp9OnTJVhdxcUdMgAABqAgAwBgAAoyAAAGoCADAGCACtVtkZWV5TT++eef1Ry7pi6gpOxOPXJ9cs6SJUvUnB9++MGt619zzTUqs2tGrFWrlsqefvppp3HTpk3VHLsGLjtfffWVyjZv3uzWa1E6mjdvrjK7k+NcT5MaOXKkmnPkyBGV2T0J7JFHHlHZJ598cqllQrhDBgDACBRkAAAMQEEGAMAAFGQAAAxQoZq6WrVq5TS+7rrr1BxfeIways7WrVtVtnbtWpV1795dZZMmTXIaDx48WM2xOzXLTrdu3VRWs2ZNt15bXHanfrmehicicvz48VJdBy7tvvvuU5ld86rrIxNXr17t1vUffvhhlTVu3Fhl//nPf9y6XkXGHTIAAAagIAMAYAAKMgAABqAgAwBgAIflejzLxSYa0uw0evRolZ09e1Zl7733nspmzJjhNLZ7/KLdCTbx8fEqs3vkWEXm5jbyGFP2o7tiY2NVtmDBAqdxYGBgWS3nkjIyMlT2xRdfqOzVV19VWXJycqmsqajYj/81f/58lfXt21dljRo1chq3bt1azXnjjTdUFhkZqTK777/333+/yvLz81Xmi9zdj9whAwBgAAoyAAAGoCADAGAACjIAAAbwuqaugwcPquzyyy8v1rXsHm0XGhparGtVdDTRFN0tt9ziNF65cmWxrzVt2jSV/fLLL269Ni8vz2k8derUYq/DFOzH/3rooYdUZte8euDAAadxSEiImmPXzDpx4kSVTZ48WWW5ubmXXKcvo6kLAAAvQkEGAMAAFGQAAAxAQQYAwABe9/jF8PBwlQ0ZMkRlEyZMUJlrk8vNN9/suYUBReT6eDs/P79yWgl82eHDh92ad+WVVzqNd+3apebYnba1Y8eOYq0LGnfIAAAYgIIMAIABKMgAABjA6w4GcVerVq1UdvLkSadxenp6WS3H53EQA0zCfoRJOBgEAAAvQkEGAMAAFGQAAAxAQQYAwAA+29SFskUTDUzCfoRJaOoCAMCLUJABADAABRkAAANQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADAABRkAAAO4fVIXAAAoPdwhAwBgAAoyAAAGoCADAGAACjIAAAagIIvIoUOHxOFwyMsvv+yxa27cuFEcDods3LjRY9dExcB+hEnYj2XHawvy/PnzxeFwyLZt28p7KaUmPT1d+vfvL8HBwVKzZk259dZb5fvvvy/vZcFGRdiP69evl549e0q9evUkODhYIiMj5d133y3vZcGGr+/H5s2bi8PhsP3VokWL8l5esVUu7wXA3qlTp6Rnz55y4sQJ+fvf/y7+/v4ydepU6dGjh+zYsUPq1q1b3ktEBbJq1SqJiYmRa6+9VhITE8XhcMjSpUslLi5OMjMzZcSIEeW9RFQg06ZNk1OnTjllP/zwg4wdO1ZuvPHGclpVyVGQDTVr1ixJS0uT1NRU6dSpk4iI9OnTR9q1aydTpkyRiRMnlvMKUZHMmDFDGjduLBs2bJCAgAARERk6dKi0atVK5s+fT0FGmYqJiVHZP/7xDxERuffee8t4NZ7jtf9k7Y6cnBwZP368XHPNNVKrVi2pVq2adO/eXZKTky/6mqlTp0pISIgEBQVJjx49ZPfu3WrO3r17JTY2VurUqSOBgYESEREhq1atKnQ9p0+flr1790pmZmahc5OSkqRTp04FxVhEpFWrVtK7d29ZunRpoa+Hebx5P2ZnZ0vt2rULirGISOXKlaVevXoSFBRU6OthHm/ej3bee+89ueKKK6Rr167Fer0JfLogZ2dny5tvvilRUVEyefJkSUxMlIyMDImOjpYdO3ao+e+8845Mnz5dhg0bJk8//bTs3r1bevXqJUeOHCmYs2fPHunSpYt8++23MmbMGJkyZYpUq1ZNYmJiZMWKFZdcT2pqqrRu3VpmzJhxyXn5+fmya9cuiYiIUL8XGRkpBw4ckJMnT7r3RYAxvHU/iohERUXJnj17ZNy4cbJ//345cOCATJgwQbZt2yajR48u8tcC5c+b96Orr776Sr799lu55557ivxao1he6u2337ZExNq6detF5+Tm5lrnzp1zyn777TerYcOG1pAhQwqygwcPWiJiBQUFWYcPHy7IU1JSLBGxRowYUZD17t3bCg8Pt86ePVuQ5efnW127drVatGhRkCUnJ1siYiUnJ6ssISHhkn+2jIwMS0Ss5557Tv3ezJkzLRGx9u7de8lroGz58n60LMs6deqU1b9/f8vhcFgiYomIVbVqVWvlypWFvhZlz9f3o6uRI0daImJ98803RX6tSXz6DtnPz0+qVKkiIn/cdWZlZUlubq5ERETI9u3b1fyYmBhp2rRpwTgyMlI6d+4sa9euFRGRrKws2bBhg/Tv319OnjwpmZmZkpmZKceOHZPo6GhJS0uT9PT0i64nKipKLMuSxMTES677zJkzIiJO/zx4QWBgoNMceA9v3Y8if+zFli1bSmxsrCxevFgWLlwoERERMnDgQPniiy+K+JWACbx5P/5Zfn6+vP/++9KxY0dp3bp1kV5rGp9v6lqwYIFMmTJF9u7dK+fPny/Ir7jiCjXXrl2+ZcuWBZ/Z7t+/XyzLknHjxsm4ceNs3+/o0aNOm7Y4Lnwmd+7cOfV7Z8+edZoD7+KN+1FEZPjw4fLFF1/I9u3bpVKlP/4e379/f2nbtq3Ex8dLSkpKid8DZc9b9+Of/fvf/5b09HSfaCz06YK8cOFCGTRokMTExMiTTz4pDRo0ED8/P5k0aZIcOHCgyNfLz88XEZFRo0ZJdHS07ZywsLASrVlEpE6dOhIQECC//PKL+r0LWZMmTUr8Pihb3rofc3JyZN68eTJ69OiCYiwi4u/vL3369JEZM2ZITk5Owd0WvIO37kdXixYtkkqVKsndd9/t8WuXNZ8uyElJSRIaGirLly8Xh8NRkCckJNjOT0tLU9m+ffukefPmIiISGhoqIn98I7rhhhs8v+D/V6lSJQkPD7f9of6UlBQJDQ2VGjVqlNr7o3R46348duyY5ObmSl5envq98+fPS35+vu3vwWzeuh//7Ny5c7Js2TKJioryiZsUn/8MWUTE+tMjn1NSUmTLli2281euXOn0GUdqaqqkpKRInz59RESkQYMGEhUVJXPnzrW9e83IyLjkeorS1h8bGytbt251KsrfffedbNiwQe68885CXw/zeOt+bNCggQQHB8uKFSskJyenID916pSsXr1aWrVqxUcoXshb9+OfrV27Vo4fP+7VP3v8Z15/h/zWW2/JunXrVB4fHy/9+vWT5cuXy2233SZ9+/aVgwcPypw5c6RNmzbqlBeRP/45pVu3bvLoo4/KuXPnZNq0aVK3bl2nH+uYOXOmdOvWTcLDw+Whhx6S0NBQOXLkiGzZskUOHz4sO3fuvOhaU1NTpWfPnpKQkFBo48Lf/vY3eeONN6Rv374yatQo8ff3l1deeUUaNmwoI0eOdP8LhDLli/vRz89PRo0aJWPHjpUuXbpIXFyc5OXlybx58+Tw4cOycOHCon2RUGZ8cT/+2aJFiyQgIEDuuOMOt+Ybr9z6u0voQlv/xX799NNPVn5+vjVx4kQrJCTECggIsDp27GitWbPGuv/++62QkJCCa11o63/ppZesKVOmWM2aNbMCAgKs7t27Wzt37lTvfeDAASsuLs5q1KiR5e/vbzVt2tTq16+flZSUVDDHE239P/30kxUbG2vVrFnTql69utWvXz8rLS2tuF8ylKKKsB8XLVpkRUZGWsHBwVZQUJDVuXNnp/eAOSrCfjxx4oQVGBho3X777cX9MhnHYVl/+vcKAABQLnz6M2QAALwFBRkAAANQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADCA2yd1/fmsU8BVWf84O/sRl8J+hEnc3Y/cIQMAYAAKMgAABqAgAwBgAAoyAAAGoCADAGAACjIAAAagIAMAYAAKMgAABqAgAwBgAAoyAAAGoCADAGAACjIAAAagIAMAYAAKMgAABqAgAwBgAAoyAAAGoCADAGAACjIAAAaoXN4LAHBxdevWVdl9992nsttvv11l3bt3V1lSUpLKPvvsM6fxa6+9VpQlAvAQ7pABADAABRkAAANQkAEAMAAFGQAAAzgsy7LcmuhwlPZa4MXc3EYeU1H249ixY1X27LPPquzs2bMqW7JkicoGDhyosjNnzjiNQ0JC1Jzjx49fapnGYT/CJO7uR+6QAQAwAAUZAAADUJABADAABRkAAAPQ1FXKwsLCVPbOO++U6nvefPPNKvv9999Vdv78eY+9J000JdeuXTuVffzxxypr1KiRyuwavZ577jmVffPNNyr7y1/+4jSePn26mtO8eXOVPf744yr74YcfVFYeKsJ+vP7661Vm99/u559/VtmyZctUNm/ePM8sDApNXQAAeBEKMgAABqAgAwBgAAoyAAAGoKnLgzp06KCyTz75RGV2j9QrbXanNj322GNO48zMzGJfvyI00ZS2CRMmqOzvf/+7ytLS0lQWHh6uMrumvWHDhqls4sSJTuPq1atfcp0XJCcnq+yGG25w67WlrSLsx/T0dJXZNfy567fffnMaf/DBB2qOXTOY6+tERL788stir8MX0dQFAIAXoSADAGAACjIAAAagIAMAYIDK5b0AX2J3clF5NHDZGTBggMreeOMNp7Fdkw7KTmBgoFvz8vLyVObuqWszZ85UmWtD0uTJk91a2/r16916T5SOTZs2qax9+/Yq8/PzU1loaKjK6tSp4zR++OGH1Ry7zG4/Hjx4UGV2jW9bt25VmWvT4o033qjmnDx5UmVPPvmkynbt2qUyk3GHDACAASjIAAAYgIIMAIABKMgAABjAZ5u6qlWrprL69es7je2asOyaBexOrGnbtq3K+vXr59bazp496zS2ezTi/v37Vfbpp5+qbPTo0SqrUqWKW+uAWTIyMlRm1wjj6VOhFi9e7DS2OzEsKChIZcuXL/foOlA0do2adipX1t/mr776apWNGjXKaWx36lqtWrVUZtc0ZvfYWbt9e+WVV6qsuFzXLyISFxfnseuXBe6QAQAwAAUZAAADUJABADAABRkAAAP4RFNXkyZNVDZ37lyV3XzzzcW6vt0j8Ny1evVqla1Zs8Zp/Oabbxb7+nbNGcX9c8I8do9t8/SjBV955RWncc2aNdWcpKQkle3bt8+j60DpyM3NVVlqaqrK+vfv7zQODg5Wc7p06aKy6OholdmdUGjX1BUbG6uy4jalmnIqYklwhwwAgAEoyAAAGICCDACAAbzuM2S7zxdcn1okInLTTTeVxXKczJ8/X2WPPfaYyuwOAikP8fHxTmOe9uQdatSooTLXJ/WIiGRlZamsXbt2KrM7AMIVh4BUPMePH1fZunXr3Mrcdd9996ksMTHRaTx+/Hi3rjVs2LBir8MU3CEDAGAACjIAAAagIAMAYAAKMgAABjC6qWvEiBEqe/TRR1XmySeG2LH7wfoFCxaozO5pI55s4LJr5gkMDCz29ewOe0D5WbRokcomTZqksqZNm6rs2WefVZndgTOuh9KIiDRq1Mhp/N5776k5y5YtUxlQGq655hqnsd1BOHZP4Dt06FBpLanMcIcMAIABKMgAABiAggwAgAEoyAAAGMCYpq7GjRurLCoqSmWl3cBlx66B6+GHHy7zdTz44IMq69WrV7Gvl56eXpLlwMOOHj2qsk2bNqmsW7duKrvjjjtU1r59e5XZPRnthx9+cBonJCSoOefPn1cZUFJ2J3X9z//8T6Gv+/LLL4v9nnanPV522WUq+/7774v9HsXFHTIAAAagIAMAYAAKMgAABqAgAwBgAGOautq0aaOyfv36lfk67B6haHcCV2m7/PLLVXbXXXeV+TpQduwap1555RWVde/eXWWup22J2DdK5uTkqCwuLs5pXB7NLCi6iRMnqiwiIkJldif82Z1qtWXLFqex3QlZJeFwOFR22223qczf37/Qa/Xs2VNl586dU5lds2PNmjVVVq9ePZU1a9as0HV4GnfIAAAYgIIMAIABKMgAABiAggwAgAEclpuf3Nt9IO9JGRkZKqtTp45H3+Prr792Gvft21fNOX78uMo8+QhFd9k1tH344YfFvt706dNVNmbMGKexXVOEuzzdAFKY0t6PprD7f2D//v0qq1WrlsrsvkY33nijytavX1/M1ZmrIuxHT3/PdP0zlEVTlyff4+eff1ZZcnKyylzrgIjI6tWrVbZ3717PLEzc/3NyhwwAgAEoyAAAGICCDACAASjIAAAYwJiTuuxOSsnPzy/29TZv3qyyu+++22lcXo8fbNu2rdPY7uSll156qdjX//XXX1W2YcMGlZWkiQtl495771WZXQOXu3766aeSLAcGcT1ZS8T90w3tvvcdO3bMaexuI5JdM5XdiWENGjRw63qu7L6Xv/rqqypLSkoq1vVNwh0yAAAGoCADAGAACjIAAAagIAMAYABjmro8rXr16iqrXLl0/7jDhw9X2fXXX6+y0NBQp3HHjh09uo7Bgwer7JNPPvHoe8DzoqOjVfbiiy+qzO4RillZWSqze/ziDTfcoLLvvvvO3SXCIHaPYw0MDHTrtXZ7yO7xn+6waw61e+Tjjh07VNa8eXOVpaamOo3tHrWYm5vr/gK9CHfIAAAYgIIMAIABKMgAABjAmM+Q7Z5g4/pZa1G0b99eZa6fYeTl5RX7+naqVq2qsoCAAI9d/+jRoyqz+6H8f//73x57T5Sehg0bOo1feOEFNadKlSoqe+SRR1T27bffquyzzz5T2cSJE1Xm+qSbH3/8US8Wxjl9+rRbWXmIj49X2RVXXKEyu8NHli5d6jT21c+L7XCHDACAASjIAAAYgIIMAIABKMgAABjAmKYuu0YVTx9mUbNmTY9erzTNmjVLZZ9++qnKVq1aVRbLQSl4/vnnncZ2jYivv/66yt544w23ru9wOFRmd2BDhw4dnMY0daEounXrprInnnjCrdcuX75cZbNnzy7xmrwVd8gAABiAggwAgAEoyAAAGICCDACAAYxp6tq5c6fKXE9sERHp379/WSynVK1du9ZpbNfAtX79epUV92ksKH/33nuvyu644w6ncXp6upozatQot65v95Qfu1OQ7DLAXXaNsW+99ZbKatWqpTK7/f3000+r7OzZs8VcnffjDhkAAANQkAEAMAAFGQAAA1CQAQAwgDFNXZmZmSq77777VGZ3oteaNWtKZU1F9eKLL6rM7nQt18c+0qzl++z2rWvji11zzODBg1XWq1cvlV1++eUlWB3gnmHDhqksLCzMrdfaNXDZPXa3IuMOGQAAA1CQAQAwAAUZAAADUJABADCAw3Lz6B67R7kBF5T1CVDeth8///xzlV133XVO4zNnzqg5didwuevXX39Vmd2j7VybEXNycor9nqZgP5ac3aM/Bw0apDI/Pz+Vbdu2TWXXX3+9ys6dO1e8xXkZd/cjd8gAABiAggwAgAEoyAAAGICCDACAAWjqgkfQRHNp7dq1U1nfvn2dxrfddpua06lTJ5Vt3bpVZV9//bXKxo4dq7IjR45ccp2+gv1YdD169HAa250yaNfA5XryoIj9Xv7oo49KsDrvRlMXAABehIIMAIABKMgAABiAggwAgAFo6oJH0EQDk7Afi65Pnz5OY3cfa5uYmKiyCRMmeGJJPoOmLgAAvAgFGQAAA1CQAQAwQOXyXgAAoPy5PqHp5MmTak5ycrLK3nvvvVJbU0XDHTIAAAagIAMAYAAKMgAABqAgAwBgAA4GgUdwEANMwn6ESTgYBAAAL0JBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAM4PZJXQAAoPRwhwwAgAEoyAAAGICCDACAASjIAAAYgIIMAIABKMgAABiAggwAgAEoyAAAGICCDACAAf4PkLEsNK/INnsAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 600x600 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# plot a 3x3 grid of MNIST digits\n",
    "idxs = np.random.randint(0, len(X_train), size=(3, 3))\n",
    "fig, axes = plt.subplots(3, 3, figsize=(3 * 2, 3 * 2))\n",
    "\n",
    "for i in range(3):\n",
    "  for j in range(3):\n",
    "    axes[i, j].imshow(X_train[idxs[i, j]], cmap=\"gray\")\n",
    "    axes[i, j].axis(\"off\")\n",
    "    axes[i, j].set_title(f\"Label: {y_train[idxs[i, j]]}\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the Model\n",
    "\n",
    "To create a convolutional neural network using NNX define a `nnx.Module` subclass. We define the model by subclassing `nnx.Module` and defining a `forward` method that returns the model output. Like in PyTorch, the `__init__` method instantiates all the modules that will be used in the model. The `__call__` in this case\n",
    "will define the forward computation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1, 10)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from flax.experimental import nnx\n",
    "\n",
    "\n",
    "class CNN(nnx.Module):\n",
    "\n",
    "  def __init__(self, *, rngs: nnx.Rngs):\n",
    "    self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n",
    "    self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n",
    "    self.linear1 = nnx.Linear(7 * 7 * 64, 256, rngs=rngs)\n",
    "    self.linear2 = nnx.Linear(256, 10, rngs=rngs)\n",
    "    self.num_calls = nnx.var(\"counts\", 0)\n",
    "\n",
    "  def __call__(self, x: jax.Array) -> jax.Array:\n",
    "    self.num_calls += 1\n",
    "    x = self.conv1(x)\n",
    "    x = nnx.relu(x)\n",
    "    x = nnx.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n",
    "    x = self.conv2(x)\n",
    "    x = nnx.relu(x)\n",
    "    x = nnx.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n",
    "    x = x.reshape((x.shape[0], -1))  # flatten\n",
    "    x = self.linear1(x)\n",
    "    x = nnx.relu(x)\n",
    "    x = self.linear2(x)\n",
    "    return x\n",
    "\n",
    "\n",
    "model = CNN(rngs=nnx.Rngs(0))\n",
    "\n",
    "y = model(X_train[:1])\n",
    "y.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One notable difference with other frameworks is that `__init__`, by convention, accepts a `rngs: nnx.Rngs` keyword-only argument. This object is passed around to generate PRNG keys as random state is explicit in JAX.\n",
    "\n",
    "One of the nice things about NNX is that Module contain their own state, are fully inspectable, and you can run them eargerly. For example, we can easily check out the kernel shape of the first `Conv` layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3, 3, 1, 32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.conv1.kernel.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also view the entire `State` of the model using the `.filter()` method. TODO: talk about collections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'conv1/bias': Variable(\n",
       "    collection='params',\n",
       "    value=(32,)\n",
       "  ),\n",
       "  'conv1/kernel': Variable(\n",
       "    collection='params',\n",
       "    value=(3, 3, 1, 32)\n",
       "  ),\n",
       "  'conv2/bias': Variable(\n",
       "    collection='params',\n",
       "    value=(64,)\n",
       "  ),\n",
       "  'conv2/kernel': Variable(\n",
       "    collection='params',\n",
       "    value=(3, 3, 32, 64)\n",
       "  ),\n",
       "  'linear1/bias': Variable(\n",
       "    collection='params',\n",
       "    value=(256,)\n",
       "  ),\n",
       "  'linear1/kernel': Variable(\n",
       "    collection='params',\n",
       "    value=(3136, 256)\n",
       "  ),\n",
       "  'linear2/bias': Variable(\n",
       "    collection='params',\n",
       "    value=(10,)\n",
       "  ),\n",
       "  'linear2/kernel': Variable(\n",
       "    collection='params',\n",
       "    value=(256, 10)\n",
       "  )\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.tree_map(jnp.shape, model.extract(nnx.Param))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training in eager mode\n",
    "\n",
    "For pedagogical purposes, we first train the model in eager mode. This will be uselful to take a look at some of NNX's features, its be more approachable for new users, and great for debugging, but it is not the recommended way to train models in JAX.\n",
    "\n",
    "Here we will run a simple `for` loop for just 10 iterations, at each step we will sample a batch of data, define a `loss_fn` to compute the loss, and use `nnx.value_and_grad` to compute the gradients of the loss with respect to the model parameters. Using the gradients we will update the parameters using stochastic gradient descent (SGD) via a simple `tree_map` operation. Finally, we will update the model's parameters using the `.update_state` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0: loss=58.7676\n",
      "Step 1: loss=80.0420\n",
      "Step 2: loss=108.3005\n",
      "Step 3: loss=26.6188\n",
      "Step 4: loss=10.7236\n",
      "Step 5: loss=4.7499\n",
      "Step 6: loss=3.9177\n",
      "Step 7: loss=2.9419\n",
      "Step 8: loss=2.4733\n",
      "Step 9: loss=1.8060\n"
     ]
    }
   ],
   "source": [
    "import optax\n",
    "\n",
    "for step in range(10):\n",
    "  idxs = np.random.randint(0, len(X_train), size=32)\n",
    "  x = jnp.array(X_train[idxs])\n",
    "  y = jnp.array(y_train[idxs])\n",
    "\n",
    "  def loss_fn(model: CNN):\n",
    "    logits = model(x)\n",
    "    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()\n",
    "\n",
    "  loss, grads = nnx.value_and_grad(loss_fn, wrt=\"params\")(model)\n",
    "  params = model.extract(\"params\")\n",
    "  params = jax.tree_map(lambda w, g: w - 0.001 * g, params, grads)\n",
    "\n",
    "  model.update(params)\n",
    "  print(f\"Step {step}: loss={loss:.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss is going down 🎉.\n",
    "\n",
    "### Training with the Functional API\n",
    "\n",
    "Now that we have a working model, lets see how to train it with `jax.jit` using NNX's Functional API. The `Module.split` method allows you to convert a Module into pytrees with functional semantics, this allows you to integrate with JAX's functional APIs like `jax.jit` and `jax.grad`.\n",
    "\n",
    "In this next example we will use the `.split` method to split the model into a `params: State` and `moduledef: ModuleDef` objects. We pass the `\"params\"` filter to check that the Module's state only contain `Variables` with the `params` collection. Having `params` and `moduledef` its pretty easy to implement a jitted `train_step` much like you would in Flax or Haiku. `ModuleDef` exposes an `apply` method which accepts some `State` and creates a function that runs the Module's `__call__` method. This function then returns the output of the Module along with the updated state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "params, moduledef = model.split(\"params\")\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def train_step(params: nnx.State, x, y):\n",
    "  def loss_fn(params):\n",
    "    logits, _updates = moduledef.apply(params)(x)\n",
    "    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()\n",
    "\n",
    "  loss, grads = jax.value_and_grad(loss_fn)(params)\n",
    "  params = jax.tree_map(lambda w, g: w - 0.001 * g, params, grads)\n",
    "\n",
    "  return loss, params"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `train_step` we can run a few more iterations and see that the loss is still going down, however, this time execution should be much faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0: loss=1.4396\n",
      "Step 1: loss=1.4127\n",
      "Step 2: loss=1.8718\n",
      "Step 3: loss=1.7080\n",
      "Step 4: loss=1.7984\n",
      "Step 5: loss=1.0350\n",
      "Step 6: loss=1.2076\n",
      "Step 7: loss=0.9081\n",
      "Step 8: loss=0.8217\n",
      "Step 9: loss=0.6687\n"
     ]
    }
   ],
   "source": [
    "for step in range(10):\n",
    "  idxs = np.random.randint(0, len(X_train), size=32)\n",
    "  x = jnp.array(X_train[idxs])\n",
    "  y = jnp.array(y_train[idxs])\n",
    "\n",
    "  loss, params = train_step(params, x, y)\n",
    "  print(f\"Step {step}: loss={loss:.4f}\")\n",
    "\n",
    "model.update(params)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Realistic Training using TrainState\n",
    "\n",
    "For real training scenarios, we recommend using `TrainState` to manage the state of your training loop. `TrainState` manages the `params` of your network along with other types of state, and uses `optax` to update the parameters according to the gradients.\n",
    "\n",
    "Next, we will define a `train_step` function that accepts a `TrainState` and a batch of data, and returns a new `TrainState` with updated parameters. The `apply_gradients` method will return a new `state` with the updated parameters. Flax users should be familiar with this API. In this case will will also define a `eval_step` function that will be used to evaluate the model on the test set and return some metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = nnx.TrainState(\n",
    "    apply_fn=moduledef.apply,\n",
    "    params=params,\n",
    "    tx=optax.adam(0.001),\n",
    ")\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def train_step(state: nnx.TrainState, x, y):\n",
    "  def loss_fn(params):\n",
    "    logits, _updates = state.apply_fn(params)(x)\n",
    "    return optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()\n",
    "\n",
    "  grads = jax.grad(loss_fn)(state.params)\n",
    "\n",
    "  state = state.apply_gradients(grads=grads)\n",
    "\n",
    "  return state\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def eval_step(state: nnx.TrainState, x, y):\n",
    "  logits, _updates = state.apply_fn(state.params)(x)\n",
    "  metrics = {\n",
    "      'accuracy': jnp.mean(jnp.argmax(logits, axis=-1) == y),\n",
    "      'loss': optax.softmax_cross_entropy_with_integer_labels(logits, y).mean(),\n",
    "  }\n",
    "  return metrics"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now lets create a simple training loop that runs for 1000 iterations and prints the metrics every 100 steps. At the end of training we will compute the final metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step 0: {'accuracy': Array(0.63119996, dtype=float32), 'loss': Array(1.1837534, dtype=float32)}\n",
      "Step 100: {'accuracy': Array(0.9492, dtype=float32), 'loss': Array(0.16359854, dtype=float32)}\n",
      "Step 200: {'accuracy': Array(0.9564, dtype=float32), 'loss': Array(0.14198248, dtype=float32)}\n",
      "Step 300: {'accuracy': Array(0.96279997, dtype=float32), 'loss': Array(0.12757339, dtype=float32)}\n",
      "Step 400: {'accuracy': Array(0.97169995, dtype=float32), 'loss': Array(0.09900841, dtype=float32)}\n",
      "Step 500: {'accuracy': Array(0.96889997, dtype=float32), 'loss': Array(0.10143881, dtype=float32)}\n",
      "Step 600: {'accuracy': Array(0.9745, dtype=float32), 'loss': Array(0.08513925, dtype=float32)}\n",
      "Step 700: {'accuracy': Array(0.96379995, dtype=float32), 'loss': Array(0.11632324, dtype=float32)}\n",
      "Step 800: {'accuracy': Array(0.97679996, dtype=float32), 'loss': Array(0.07204168, dtype=float32)}\n",
      "Step 900: {'accuracy': Array(0.9765, dtype=float32), 'loss': Array(0.08413408, dtype=float32)}\n",
      "Final metrics: {'accuracy': Array(0.9819, dtype=float32), 'loss': Array(0.05711861, dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "total_steps = 1000\n",
    "eval_every = 100\n",
    "\n",
    "for step in range(total_steps):\n",
    "  if step % eval_every == 0:\n",
    "    metrics = eval_step(state, jnp.array(X_test), jnp.array(y_test))\n",
    "    print(f\"Step {step}: {metrics}\")\n",
    "\n",
    "  idxs = np.random.randint(0, len(X_train), size=32)\n",
    "  x = jnp.array(X_train[idxs])\n",
    "  y = jnp.array(y_train[idxs])\n",
    "\n",
    "  state = train_step(state, x, y)\n",
    "\n",
    "metrics = eval_step(state, jnp.array(X_test), jnp.array(y_test))\n",
    "print(f\"Final metrics: {metrics}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference\n",
    "\n",
    "Finally, now that we have a trained model, lets use it to make some predictions. We will update the `model` object with the trained parameters and use it to make predictions on the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAeQAAAH4CAYAAACbup4ZAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAABBzklEQVR4nO3de1hVVf7H8S83lXukqJiGaGrmJW9Zk5cUUUa8JOakZY1iTVTeffLalKWOllrpoJk2hdXgpKaMk6GOlk6ieSnJUrPMMDXGS5PiDS/A+v3hD2qztnA4HDgLeL+exz/Wh7X3XpxWfNnnLNb2UEopAQAAbuXp7gEAAAAKMgAARqAgAwBgAAoyAAAGoCADAGAACjIAAAagIAMAYAAKMgAABqAgAwBggEpXkOvXry9Dhw7Nb2/ZskU8PDxky5YtLruGh4eHvPDCCy47Hyou5iNMwnx0rzItyEuXLhUPD4/8f9WqVZPGjRvLiBEj5OTJk2U5lBJLSUkpN5Nq165d8vTTT0vbtm3Fx8dHPDw83D0kIzAfy15ubq4sXbpU+vbtK/Xq1RN/f39p3ry5zJgxQy5fvuzu4bkV89F9FixYIE2bNpWqVavKLbfcIuPGjZOLFy+W+Ti8y/yKIjJt2jSJiIiQy5cvS2pqqixatEhSUlJk37594ufnV6Zj6dy5s2RlZUmVKlWKdVxKSoosXLjQdtJlZWWJt7dbXlpbKSkp8re//U1atmwpDRo0kO+++87dQzIK87HsXLp0SeLi4uSee+6RJ598UmrWrCmfffaZTJ06VT7++GP55JNPKv0vjMzHsjVx4kSZPXu2DBgwQEaPHi0HDhyQhIQE2b9/v2zYsKFsB6PKUGJiohIRtXv3bks+btw4JSJq2bJlNzz2woULLhlDeHi4GjJkSInPM3z4cFXGL5/TTpw4oS5duqSUKl/jLm3Mx7J35coVtW3bNi1/8cUXlYiojRs3umFUZmA+lr2MjAzl7e2tHn30UUuekJCgRET961//KtPxGPEZcmRkpIiIpKeni4jI0KFDJSAgQA4fPiwxMTESGBgogwcPFpHrb3nNmzdPmjVrJtWqVZNatWpJfHy8nDlzxnJOpZTMmDFD6tatK35+ftK1a1fZv3+/du0bfUayc+dOiYmJkZCQEPH395eWLVvK/Pnz88e3cOFCERHLW0x57D4jSUtLk549e0pQUJAEBARIt27dZMeOHZY+eW9Zbdu2TcaNGyehoaHi7+8vsbGxcvr0aUvfzMxMOXjwoGRmZhb5+taqVUt8fX2L7IfrmI/XlcZ8rFKlitx7771aHhsbKyIi33zzTaHHV0bMx+tKYz5+9tlnkp2dLYMGDbLkee3333+/0ONdzYj3DQ4fPiwiItWrV8/PsrOzJTo6Wjp27Chz587Nf6smPj5eli5dKnFxcTJq1ChJT0+XBQsWSFpammzbtk18fHxEROT555+XGTNmSExMjMTExMiePXukR48ecvXq1SLHs3HjRundu7eEhYXJ6NGjpXbt2vLNN9/I2rVrZfTo0RIfHy8ZGRmyceNGee+994o83/79+6VTp04SFBQkEyZMEB8fH1m8eLF06dJF/vOf/8jdd99t6T9y5EgJCQmRqVOnypEjR2TevHkyYsQIWb58eX6f5ORkiYuLk8TERMsiDJQc87Hs5+OJEydERKRGjRrFPraiYz6W3ny8cuWKiIh2w5L3en7xxRdFjt+lyvJ2PO8tmU2bNqnTp0+rY8eOqffff19Vr15d+fr6quPHjyullBoyZIgSETVp0iTL8Vu3blUiopKSkiz5+vXrLfmpU6dUlSpVVK9evVRubm5+vylTpigRsbwls3nzZiUiavPmzUoppbKzs1VERIQKDw9XZ86csVznt+cq7C0ZEVFTp07Nb/fr109VqVJFHT58OD/LyMhQgYGBqnPnztrrExUVZbnW2LFjlZeXlzp79qzWNzEx0XYMN1Je3koqC8xH98/HPFFRUSooKEj7HisT5mPZz8cvvvhCiYiaPn26Jc97zQICAgo93tXc8pZ1VFSUhIaGSr169WTQoEESEBAgycnJcsstt1j6PfXUU5b2ypUrJTg4WLp37y4///xz/r+2bdtKQECAbN68WURENm3aJFevXpWRI0da3ioZM2ZMkWNLS0uT9PR0GTNmjNx0002Wrzmz2CQnJ0f+/e9/S79+/aRBgwb5eVhYmDz88MOSmpoq586dsxzzxBNPWK7VqVMnycnJkR9//DE/Gzp0qCiluDt2Aeaje+fjzJkzZdOmTfLSSy9p32NlxHwsu/nYpk0bufvuu+Xll1+WxMREOXLkiKxbt07i4+PFx8dHsrKyiv09lYRb3rJeuHChNG7cWLy9vaVWrVrSpEkT8fS0/m7g7e0tdevWtWSHDh2SzMxMqVmzpu15T506JSKS/x+mUaNGlq+HhoZKSEhIoWPLe3uoefPmjn9DhTh9+rRcunRJmjRpon2tadOmkpubK8eOHZNmzZrl57feequlX96YC34OBNdgPl7njvm4fPly+fOf/yyPPfaYVmAqK+bjdWU1H1etWiUDBw6UYcOGiYiIl5eXjBs3Tv7zn//It99+69Q5neWWgty+fXtp165doX2qVq2qTcLc3FypWbOmJCUl2R4TGhrqsjG6k5eXl22ulCrjkVQOzMfCldZ83Lhxo/zxj3+UXr16yRtvvFGic1UkzMfCuXo+3nLLLZKamiqHDh2SEydOSKNGjaR27dpSp04dady4cUmGWmxGLOpyVMOGDWXTpk3SoUOHQlcNh4eHi8j13xh/+zbI6dOni/wtqmHDhiIism/fPomKirphP0ffngkNDRU/Pz/b37QOHjwonp6eUq9ePYfOBbMwH523c+dOiY2NlXbt2smKFSuM+rvU8or5WDKNGjXKf9fgwIED8t///rfMPxI04s+eHPXggw9KTk6OTJ8+Xftadna2nD17VkSufwbj4+MjCQkJlt+a5s2bV+Q12rRpIxERETJv3rz88+X57bn8/f1FRLQ+BXl5eUmPHj1kzZo1cuTIkfz85MmTsmzZMunYsaMEBQUVOa6CivNnTygdzMdfFWc+fvPNN9KrVy+pX7++rF27lj/JcxHm469K8vMxNzdXJkyYIH5+fvLkk08W+/iSKFe/lt53330SHx8vs2bNki+//FJ69OghPj4+cujQIVm5cqXMnz9fBgwYIKGhofLMM8/IrFmzpHfv3hITEyNpaWmybt26Iv+swtPTUxYtWiR9+vSRVq1aSVxcnISFhcnBgwctO7e0bdtWRERGjRol0dHR4uXlpf0tW54ZM2bIxo0bpWPHjvL000+Lt7e3LF68WK5cuSKzZ8926rUozp+Z/Pjjj/l/fvD555/nj0nk+m/Ljz76qFNjqOyYj79ydD6eP39eoqOj5cyZMzJ+/Hj56KOPLF9v2LCh/O53v3NqDJUd8/FXxfn5OHr0aLl8+bK0atVKrl27JsuWLZNdu3bJO++8o31eXerKckn3jXaiKWjIkCHK39//hl9fsmSJatu2rfL19VWBgYGqRYsWasKECSojIyO/T05OjnrxxRdVWFiY8vX1VV26dFH79u3TdqIpuKw/T2pqqurevbsKDAxU/v7+qmXLliohISH/69nZ2WrkyJEqNDRUeXh4WJb4S4Fl/UoptWfPHhUdHa0CAgKUn5+f6tq1q9q+fbtDr4/dGIvzZyZ5x9v9u++++4o8vqJiPpb9fExPT7/hXJQCf3JT2TAf3fPzMTExUd15553K399fBQYGqm7duqlPPvmkyONKg4dSrBQCAMDdytVnyAAAVFQUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAO79TlzKO1UHmU9Z+zMx9RGOYjTOLofOQOGQAAA1CQAQAwAAUZAAADUJABADAABRkAAANQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADAABRkAAANQkAEAMAAFGQAAA1CQAQAwgMOPXwTgvOjoaC0bP368pR0ZGenSa9o9EnDNmjVatn37dkt73rx5Wp+rV6+6bFwA7HGHDACAASjIAAAYgIIMAIABPJRSyqGONp9HAXkcnEYuY/J8rF+/vpbt379fy6pVq1YGoym+lJQULXv11Ve1bPPmzWUxHKcwH2ESR+cjd8gAABiAggwAgAEoyAAAGICCDACAAVjUBZdgEc2vGjZsqGXfffedG0biOufPn9eyyZMna9nq1au17OTJk6UypsIwH2ESFnUBAFCOUJABADAABRkAAANQkAEAMACLugpo27atlv3lL3/RMn9/fy2bPn26lv373/92zcAMxyKaX/n6+mrZzJkztSwzM9PS3rRpk9bnz3/+s5YlJiZqWUxMjJY1a9ZMy1q3bq1lrmS3y1efPn1K9Zp2mI8wCYu6AAAoRyjIAAAYgIIMAIABKMgAABiARV0FxMXFadmbb77p0LHZ2dla1qNHDy379NNPiz8ww7GIxjx2i8vsFli9/vrrlnZISIjT1/z++++1rF27dlpmt/OXKzEfCzdp0qQis6CgIIfO9dprr2nZ7NmztczRHdsK/sxs06aN1uell15y6FymYFEXAADlCAUZAAADUJABADAABRkAAAOwqKuA2267TcvWr1+vZfXr19cyu9dow4YNWma3q1J5xyKa8qvgohm7HcOCg4OdPv9jjz2mZUuXLnX6fI5gPv5q8ODBWvb3v/9dyy5cuGBpZ2VlaX3sdii0Wzz41VdfadnEiRO1zM/PT8vefffdIq/ZqlUrh65pChZ1AQBQjlCQAQAwAAUZAAADUJABADAAi7ocMHLkSC2z253G7jXKyMjQsnr16rlmYAZhEU3F8eSTT2rZwoULnT7f8uXLtezhhx92+nyOYD7+auvWrVrWoUMHLSu4u5bdbl533HGHlr3yyitaFh0drWVXr17Vsl9++UXLateurWUF/elPf9Kyt956q8jj3IVFXQAAlCMUZAAADEBBBgDAABRkAAAM4O3uAQCo2CIjI909BDjAkd3YDhw4oGWxsbFaZrcw8NVXX9UyRxZwVSbcIQMAYAAKMgAABqAgAwBgAD5DBoAK7NChQ1pmtzHIsGHDLO2kpCStT2pqqpZdvnxZy3bv3l2cIeL/cYcMAIABKMgAABiAggwAgAEoyAAAGIBFXQBK1fz58909hErtpZde0rJevXppWc2aNS3tTz/9VOszbtw4LVu7dq2WlfbTr06ePFmq53cX7pABADAABRkAAANQkAEAMAAFGQAAA7CoC4DF448/7vSxdrs27d+/vyTDQQl99913WtajRw8t+/jjjy3tm2++Wetj98SmqVOnatnnn39enCEWym7HMLuFZBUBd8gAABiAggwAgAEoyAAAGICCDACAAVjUVcoWLFjg7iHACf7+/loWEhKiZfHx8VrWoEGDUhlTYVJSUrQsMzNTy8LDw7Xs3nvvtbRbtGjh9Dh++OEHLfvXv/7l9PlQOvbu3atljz76qKX91ltvaX1q166tZcHBwVrWrVs3p8d2/vx5S/svf/mL0+cqb7hDBgDAABRkAAAMQEEGAMAAFGQAAAzAoq5SdvHiRXcPAUWoWrWqlv3973/Xsr59+5bFcJwyaNAgdw9BROwX/XTo0EHLdu7cqWXZ2dmlMiY4Zt26dZZ2kyZNtD5jxozRsoceekjL/Pz8tOzWW291aBw//fSTpX3w4EGHjqsIuEMGAMAAFGQAAAxAQQYAwAAUZAAADMCiLgd4eHg4lHl66r/f2PWDWbKysrRMKeWGkZR/do/s+/TTT7Vs4sSJWjZ37txSGROcU3DHLBGR6dOnO5TZLeTbunWrQ9e1e1xkZcEdMgAABqAgAwBgAAoyAAAGoCADAGAAFnU5wG6Bj12Wk5OjZRcuXCiVMcFc586d07J//OMfTp3rkUce0TK7R0MCJinJYtbt27e7cCTlC3fIAAAYgIIMAIABKMgAABiAz5BdyO7JTomJiW4YCcrKl19+qWX9+vXTsmPHjjl1/q+++krLWrdurWWPP/64U+e388MPP2hZWlqalkVGRmpZSEiIy8aB8qskG+v079/f0p49e3ZJh1NucIcMAIABKMgAABiAggwAgAEoyAAAGIBFXUAJ3HbbbVq2YsUKl53fbgGXj4+Py84vInL06FFL+/XXX9f6vPbaa1rWp08fLRs3bpxD1zx+/LiDo0NlExQU5O4huA13yAAAGICCDACAASjIAAAYgIIMAIABWNTlQsuXL3f3EOCEb775RssaNWqkZV5eXloWEBCgZe3bt3fNwErBTz/9pGW///3vLe1vv/3WoXN9+OGHDmVAcezevdvdQ3Ab7pABADAABRkAAANQkAEAMAAFGQAAA7CoywGrVq3SspkzZ2qZh4dHWQwHLtasWTMtmzt3rpY99thjWuaOXYWuXbumZWfPntWyf/zjH1q2ePFiLXN0ERdQFvbt2+fuIbgNd8gAABiAggwAgAEoyAAAGICCDACAAVjU5YCMjAwty83N1bK6deuWxXBQBp555hktW7RokZZt2LBByyIiIpy6ZmpqqpbZ7Xz1448/atnKlSuduiZQGkqywLUyL47lDhkAAANQkAEAMAAFGQAAA1CQAQAwgIdSSjnUsRJ/0G4nMzPToX7BwcGlPBIzODiNXIb5iMIwH92rQ4cOWrZ161aHjt21a5elfc8997hkTO7k6HzkDhkAAANQkAEAMAAFGQAAA1CQAQAwADt1Oenzzz/Xsnbt2rlhJABQcbjjkaam4A4ZAAADUJABADAABRkAAAPwGbKTpk+frmUTJ050w0gAoOJwdNOliog7ZAAADEBBBgDAABRkAAAMQEEGAMAAPO0JLsHTdWAS5qN73XbbbVqWlJSkZXfddZeWxcbGWtpr1qxx3cDchKc9AQBQjlCQAQAwAAUZAAADUJABADAAi7rgEiyigUmYjzAJi7oAAChHKMgAABiAggwAgAEoyAAAGMDhRV0AAKD0cIcMAIABKMgAABiAggwAgAEoyAAAGKDSFeT69evL0KFD89tbtmwRDw8P2bJli8uu4eHhIS+88ILLzoeKi/kIkzAf3atMC/LSpUvFw8Mj/1+1atWkcePGMmLECDl58mRZDqXEUlJSytWkys3NlUWLFkmrVq3E19dXqlevLpGRkbJ37153D81tmI/ud+3aNbnjjjvEw8ND5s6d6+7huBXz0X1M+fnoXaZX+3/Tpk2TiIgIuXz5sqSmpsqiRYskJSVF9u3bJ35+fmU6ls6dO0tWVpZUqVKlWMelpKTIwoULbSddVlaWeHu75aW9oWHDhklSUpL88Y9/lBEjRsjFixclLS1NTp065e6huR3z0X0SEhLk6NGj7h6GUZiPZc+Un49ueVV69uwp7dq1ExGRxx9/XKpXry6vvvqqrFmzRh566CHbYy5evCj+/v4uH4unp6dUq1bNped09flKasWKFfLOO+/I6tWrJTY21t3DMQ7z0T1OnTol06ZNk4kTJ8rzzz/v7uEYg/lYtkz6+WjEZ8iRkZEiIpKeni4iIkOHDpWAgAA5fPiwxMTESGBgoAwePFhErr+1MG/ePGnWrJlUq1ZNatWqJfHx8XLmzBnLOZVSMmPGDKlbt674+flJ165dZf/+/dq1b/QZyc6dOyUmJkZCQkLE399fWrZsKfPnz88f38KFC0VELG8x5bH7jCQtLU169uwpQUFBEhAQIN26dZMdO3ZY+uS9ZbVt2zYZN26chIaGir+/v8TGxsrp06ctfTMzM+XgwYOSmZlZ5Ov76quvSvv27SU2NlZyc3Pl4sWLRR5TmTEfryut+Zhn0qRJ0qRJE3nkkUccPqYyYj5eVxl+PhpRkA8fPiwiItWrV8/PsrOzJTo6WmrWrClz586VBx54QERE4uPjZfz48dKhQweZP3++xMXFSVJSkkRHR8u1a9fyj3/++eflueeekzvvvFPmzJkjDRo0kB49ejj0Ym/cuFE6d+4sBw4ckNGjR8srr7wiXbt2lbVr1+aPoXv37iIi8t577+X/u5H9+/dLp06dZO/evTJhwgR57rnnJD09Xbp06SI7d+7U+o8cOVL27t0rU6dOlaeeeko+/PBDGTFihKVPcnKyNG3aVJKTkwv9Xs6dOye7du2Su+66S6ZMmSLBwcESEBAgDRo0kBUrVhT5WlRGzEcrV87HPLt27ZJ33nlH5s2bx6MLi8B8tKrQPx9VGUpMTFQiojZt2qROnz6tjh07pt5//31VvXp15evrq44fP66UUmrIkCFKRNSkSZMsx2/dulWJiEpKSrLk69evt+SnTp1SVapUUb169VK5ubn5/aZMmaJERA0ZMiQ/27x5sxIRtXnzZqWUUtnZ2SoiIkKFh4erM2fOWK7z23MNHz5c3ejlExE1derU/Ha/fv1UlSpV1OHDh/OzjIwMFRgYqDp37qy9PlFRUZZrjR07Vnl5eamzZ89qfRMTE23HkGfPnj1KRFT16tVVrVq11Ouvv66SkpJU+/btlYeHh1q3bl2hx1dkzMeyn495427fvr166KGHlFJKpaenKxFRc+bMKfLYioz5yM9HtxTkgv/Cw8PV+vXr8/vlTbgff/zRcvyoUaNUcHCwOnXqlDp9+rTlX0BAgHr88ceVUkotW7ZMiYjlnEpdn4hFTbjdu3crEVGvvfZaod+LoxMuOztb+fn5qQcffFDrFx8frzw9PVVmZqbl9VmxYoWl3+rVq5WIqL179xY6Jjuffvpp/uu8Y8eO/Pz8+fOqRo0aqkOHDsU+Z0XBfLQqi/molFJvv/228vX1VUePHlVKUZDzMB+tKuPPR7cs6lq4cKE0btxYvL29pVatWtKkSRPx9LS+e+7t7S1169a1ZIcOHZLMzEypWbOm7XnzVsT9+OOPIiLSqFEjy9dDQ0MlJCSk0LHlvT3UvHlzx7+hQpw+fVouXbokTZo00b7WtGlTyc3NlWPHjkmzZs3y81tvvdXSL2/MBT8HcoSvr6+IiERERMjdd9+dnwcEBEifPn3k73//u2RnZxu36rEsMR+vK4v5eO7cOZk8ebKMHz9e6tWrV+zjKwPm43WV8eejW34Kt2/fPn8V4Y1UrVpVm4S5ublSs2ZNSUpKsj0mNDTUZWN0Jy8vL9tcOfFgrjp16oiISK1atbSv1axZU65duyYXL16U4ODgYp+7omA+Fs6V83Hu3Lly9epVGThwoBw5ckRERI4fPy4i13+gHjlyROrUqVPsP7OpSJiPhavIPx/L1W1Rw4YNZdOmTdKhQ4f832zshIeHi8j13xgbNGiQn58+fbrI36IaNmwoIiL79u2TqKioG/ZzdCFKaGio+Pn5ybfffqt97eDBg+Lp6Vmqdwp16tSR2rVry08//aR9LSMjQ6pVqyaBgYGldv2KjPlYfEePHpUzZ85Y7njyzJw5U2bOnClpaWnSqlWrUhtDRcV8LD7Tfj4ascraUQ8++KDk5OTI9OnTta9lZ2fL2bNnRUQkKipKfHx8JCEhwfJb07x584q8Rps2bSQiIkLmzZuXf748vz1X3t/8FexTkJeXl/To0UPWrFmTf0cgInLy5ElZtmyZdOzYUYKCgoocV0HFWdY/cOBAOXbsmGzcuDE/+/nnn2XNmjUSGRmp/aYNxzAff+XofBw1apQkJydb/i1evFhErv+5THJyskRERBT7+mA+/lZ5/flYru6Q77vvPomPj5dZs2bJl19+KT169BAfHx85dOiQrFy5UubPny8DBgyQ0NBQeeaZZ2TWrFnSu3dviYmJkbS0NFm3bp3UqFGj0Gt4enrKokWLpE+fPtKqVSuJi4uTsLAwOXjwoOzfv182bNggIiJt27YVkes/YKKjo8XLy0sGDRpke84ZM2bIxo0bpWPHjvL000+Lt7e3LF68WK5cuSKzZ8926rVITk6WuLg4SUxMtOw9a2fy5MmyYsUKeeCBB2TcuHESHBwsb7zxhly7dk1mzpzp1PXBfPwtR+djmzZtpE2bNpYs7wdxs2bNpF+/fk5dH8zH3yq3Px/LcgVZ3iq53bt3F9pvyJAhyt/f/4ZfX7JkiWrbtq3y9fVVgYGBqkWLFmrChAkqIyMjv09OTo568cUXVVhYmPL19VVdunRR+/btU+Hh4YWuIsyTmpqqunfvrgIDA5W/v79q2bKlSkhIyP96dna2GjlypAoNDVUeHh6WFYVSYFm/UteX10dHR6uAgADl5+enunbtqrZv3+7Q62M3xuL8mYlSSh0+fFjFxsaqoKAg5evrqyIjI9WuXbscOraiYj66bz7+Fqusr2M+8vPRQyknPgkHAAAuxYeHAAAYgIIMAIABKMgAABiAggwAgAEoyAAAGICCDACAASjIAAAYwOGduniIOApT1n/OznxEYZiPMImj85E7ZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAABRkAAAMQEEGAMAAFGQAAAxAQQYAwAAUZAAADEBBBgDAAA4/frEiCgoK0rKvvvpKy0aOHKllH374YamMCQDcoeAjAvfu3av1mThxopZt2LCh1MZU2XCHDACAASjIAAAYgIIMAIABPFTBDw5u1NHDo7THUuqaNWtmaSckJGh9unbtqmXvvvuulg0ZMsR1A6sAHJxGLlMR5iNKD/Ox+FauXGlp9+/fX+tz4MABLWvfvr2WZWVluW5gFYCj85E7ZAAADEBBBgDAABRkAAAMQEEGAMAAlWpjkObNm1vadgu47Jw6dao0hgNUCnPmzNGyvn37almTJk3KYji4gczMzCL7nD9/XstycnJKYziVEnfIAAAYgIIMAIABKMgAABiAggwAgAEq1aIuR/zvf//TMrsnnADQjR07VsuefvppLfPx8dGyNm3aWNp79uxx3cDgEidOnNCyq1evumEkFRN3yAAAGICCDACAASjIAAAYgIIMAIABKuyiLi8vLy17+OGHizxu586dWpabm+uSMRVHSEiIlp07d07L2CWnYqtWrZqWFXyMqIjIF198UarjsPv/adq0aVo2efJkLbN79NyyZcu0zO7Rfig7hw4dcvcQKj3ukAEAMAAFGQAAA1CQAQAwAAUZAAADVNhFXYMHD9aygo98++WXX7Q+Tz75ZKmN6UZuv/12Ldu8ebOW2S04mzFjhpZ9/vnnrhkY3O7ZZ5/VsuzsbC1z5aIuuwVcL7zwgpZNmjTJofPZLeD605/+pGWXL1926HwoHR07drS0PTw8tD7bt28vq+FUStwhAwBgAAoyAAAGoCADAGAACjIAAAaosIu6srKyiuxjtxtWYGBgaQzHomXLlpb2ggULtD61a9fWsvvvv1/LYmJitKzg4gwRkV27dhVniHCDdu3aadn48eO1bNasWaU6jsjISC2bMmWKQ8euXLlSy4YNG6Zl165dK/7A4DJ2O8D16tXL0v7LX/6i9dm0aVOpjQncIQMAYAQKMgAABqAgAwBgAAoyAAAGqBCLum655RYti4uLK/K4o0ePapnd7l2O8vTUf78ZMGCAlg0fPtzS7tSpk9PXtNtVyW4nJ5hv9OjRWlalSpVSv25sbKylvWrVKoeOs9sRbuDAgS4ZE1zHbqFqfHx8kceNGDFCyxYuXOiSMcEed8gAABiAggwAgAEoyAAAGKBCfIb82GOPaVnPnj2LPK7gZ7kiIidOnHB6HM2bN9ey5cuXO30+R+Tm5mrZnj17SvWacI0aNWpY2nYbgyiltOznn392+pp33HGHlr399ttFXvPQoUNa1qFDB6fHgbJz6dIlLXvjjTe0rODmL+vWrdP6lOTnI4rGHTIAAAagIAMAYAAKMgAABqAgAwBggAqxqMtu8w07q1evtrRTUlKcvqafn5+WJSUlOXUuu0U6O3bs0LIWLVpoWXh4uJY9+OCDWrZixQqnxgbXqFWrlpZ99NFHlnaTJk20Plu3btWyJUuWOHTNoKAgLVu8eHGR/Y4cOaL16devn5bxxKbyoW7dulr20EMPadm+ffssbbuNX0zRuHFjh/p99913pTwS1+IOGQAAA1CQAQAwAAUZAAADUJABADBAuVvUZff0G7vdh+wU3LHGbkciOz4+Plq2dOlSLbPbqctOwd1uRo4cqfX54IMPHMrq1aunZT/88IND40DpsHsC1zvvvKNlrVu3trSvXr2q9Zk+fbqWObqY6oEHHtAyu921Cv5/cODAAa1P9+7dtcxul7jytoimMnj66ae17P7779eyw4cPW9qXL18utTEVJjo62tKePXu21sdugauHh4eWJSYmatmwYcNKMLrSxR0yAAAGoCADAGAACjIAAAagIAMAYIByt6jLbhcqu0U0X3/9tZY988wzRZ7fbtHYsmXLtMxuwYyjRo8ebWnbLday8/3332uZ3UIGu13EUHbee+89LbNbFFXQm2++qWV2/83tFuS0adNGy+Lj44u8pp2YmBiHsm3btmlZ586dnbomSs/Bgwe1LDIyUsvWrl1raS9atKjUxpTHbrfEggsPAwICtD6OLsjt1auXcwNzE+6QAQAwAAUZAAADUJABADAABRkAAAOUu0VdGzdu1DK7HYPsHj138uTJIs9/6623allJFnC9//77WuboIq6C7rrrLi2zW9TVv39/Lfv000+duiaKr127dlpm99+poBEjRmjZ8OHDXTKm4ozDUdu3b3fZuVB6du7cqWV2i2MffvhhS9vVi7rsdtyy2zmu4CKuPXv2aH08PfV7Sbuf3f/85z+LMUL34w4ZAAADUJABADAABRkAAANQkAEAMEC5W9Rlt+OR3Qf86enpRZ6rTp06WpacnOzcwEQkLS1Ny15++WUtK7gIzW4nmilTpmiZ3S5Ido/i27BhQ6HjhOvUqFFDy0qys5CrjruRs2fPalnBnZzsdrmzW4i4ZcsWVw0LpSgwMFDL7B716ayqVatq2aOPPqpldjslZmdna9nChQst7QsXLmh9/vCHP2iZ3Y5kBXdFNB13yAAAGICCDACAASjIAAAYgIIMAIAByt2iri5dujjUb8eOHUX22bRpk5Y1bdrUofOnpqZq2ciRI7UsLCxMy/r27WtpjxkzRusTEhLi0DhefPFFLVu3bp1Dx6Lkfv75Zy373e9+p2U+Pj5Fnstu1yK7RyjefffdWma3WGvOnDla9u6772pZRkZGkWND+WW3c9yVK1e07LvvvnPq/Hbz0W6xlp2pU6dq2axZsyxtu0WGDRo00LITJ05o2eXLlx0ahym4QwYAwAAUZAAADEBBBgDAAB7KwZ0HXPmUmJJYvXq1lsXGxmrZpUuXtOz8+fOWdq1atZweR8Fz3YjdH+U7y27jkXvuuUfLXPlH/45y9QYWRTFlPrpSlSpVtMzu87PbbrtNy+bOnatlEydOdM3AyiHmY/H17t3b0l67dq3Wx+6JTXabHz3xxBNadu+992rZV199pWU9e/a0tI8cOaL1sXuCld04TOHofOQOGQAAA1CQAQAwAAUZAAADUJABADBAudsYZNWqVVpmt6jLz8/PocxZrlysZcduIcOzzz6rZe5YwIXSMX36dC2zW8C1bds2LSu4mQKQJzw8XMvsFljt37+/yHN9//33Wma3wNXbWy8t9erV07KCT74TETlz5oylbbfJiN2GPBUBd8gAABiAggwAgAEoyAAAGICCDACAAcrdTl12Pv74Yy2LjIx0w0gcc+jQIUv7r3/9q9ZnwYIFZTUcl2BnpOL7wx/+YGmvWLFC62P3uto9AerNN9903cAqAOZj4aKiorSs4C5cdk/Da9SokZb997//1TK7OXrs2DEtCw0N1bK4uDhL+5dfftH62I3NZOzUBQBAOUJBBgDAABRkAAAMQEEGAMAAFWJRV0hIiJbZLSoYO3aspV2zZk2nr2m3Q1ZmZqaWHThwQMsGDBhgaVeEXWdYRFM4f39/LSv4aMX69etrfVauXKlldnP77NmzTo+tImI+Fl/BXbMcfQ3tfhb+8MMPWmb3Gi1btkzLCi4uqwi7EbKoCwCAcoSCDACAASjIAAAYgIIMAIABKsSiLkfdfffdlrbdY+y8vLwcOpfd4+7sdtw6ceKEg6Mr31hEU7hXXnlFy8aMGWNpX7p0SevTs2dPLUtNTXXZuCoq5mPxFdzt7YEHHtD6BAcHa9mHH36oZf369XPZuCoCFnUBAFCOUJABADAABRkAAANQkAEAMEClWtSF0sMiml917txZy5YuXapl4eHhlvbw4cO1Pm+88YbLxlWZMB9hEhZ1AQBQjlCQAQAwAAUZAAADUJABADAAi7rgEiyi+dWGDRu0LCoqSsvefvttS/tPf/pTqY2psmE+wiQs6gIAoByhIAMAYAAKMgAABuAzZLgEn9nBJMxHmITPkAEAKEcoyAAAGICCDACAASjIAAAYgIIMAIABKMgAABiAggwAgAEoyAAAGICCDACAASjIAAAYgIIMAIABKMgAABiAggwAgAEoyAAAGMDhxy8CAIDSwx0yAAAGoCADAGAACjIAAAagIAMAYIBKV5Dr168vQ4cOzW9v2bJFPDw8ZMuWLS67hoeHh7zwwgsuOx8qLuYjTMJ8dK8yLchLly4VDw+P/H/VqlWTxo0by4gRI+TkyZNlOZQSS0lJKZeT6tq1a3LHHXeIh4eHzJ07193DcSvmo/usWLFC7rnnHrnpppukevXqct9998lHH33k7mG5FfPRPd5880257777pFatWlK1alWJiIiQuLg4OXLkSJmPxbvMrygi06ZNk4iICLl8+bKkpqbKokWLJCUlRfbt2yd+fn5lOpbOnTtLVlaWVKlSpVjHpaSkyMKFC20nXVZWlnh7u+WlLVJCQoIcPXrU3cMwCvOxbCUkJMioUaOkV69e8tJLL8nly5dl6dKl0rt3b1m1apX079/f3UN0K+Zj2UpLS5OIiAjp27evhISESHp6urz55puydu1a2bt3r9SpU6fsBqPKUGJiohIRtXv3bks+btw4JSJq2bJlNzz2woULLhlDeHi4GjJkSInPM3z4cFXGL1+JnTx5UgUHB6tp06YpEVFz5sxx95DcivnoHo0aNVJ33XWXys3Nzc8yMzNVQECA6tu3rxtH5l7MR3N8/vnnSkTUrFmzyvS6RnyGHBkZKSIi6enpIiIydOhQCQgIkMOHD0tMTIwEBgbK4MGDRUQkNzdX5s2bJ82aNZNq1apJrVq1JD4+Xs6cOWM5p1JKZsyYIXXr1hU/Pz/p2rWr7N+/X7v2jT4j2blzp8TExEhISIj4+/tLy5YtZf78+fnjW7hwoYiI5S2mPHafkaSlpUnPnj0lKChIAgICpFu3brJjxw5Ln7y3rLZt2ybjxo2T0NBQ8ff3l9jYWDl9+rSlb2Zmphw8eFAyMzMdeYlFRGTSpEnSpEkTeeSRRxw+pjJiPl5XWvPx3LlzUrNmTcsY88bh6+tb5PGVDfPxutL++fhb9evXFxGRs2fPOnW8s4x43+Dw4cMiIlK9evX8LDs7W6Kjo6Vjx44yd+7c/Ldq4uPjZenSpRIXFyejRo2S9PR0WbBggaSlpcm2bdvEx8dHRESef/55mTFjhsTExEhMTIzs2bNHevToIVevXi1yPBs3bpTevXtLWFiYjB49WmrXri3ffPONrF27VkaPHi3x8fGSkZEhGzdulPfee6/I8+3fv186deokQUFBMmHCBPHx8ZHFixdLly5d5D//+Y/cfffdlv4jR46UkJAQmTp1qhw5ckTmzZsnI0aMkOXLl+f3SU5Olri4OElMTLQswriRXbt2yTvvvCOpqamW/zmgYz6W7nzs0qWLfPDBB5KQkCB9+vSRy5cvS0JCgmRmZsro0aOLHH9lw3ws/Z+PIiL/+9//JCcnR44ePSrTpk0TEZFu3bo5dKzLlOXteN5bMps2bVKnT59Wx44dU++//76qXr268vX1VcePH1dKKTVkyBAlImrSpEmW47du3apERCUlJVny9evXW/JTp06pKlWqqF69elneFpsyZYoSEctbMps3b1YiojZv3qyUUio7O1tFRESo8PBwdebMGct1fnuuwt6SERE1derU/Ha/fv1UlSpV1OHDh/OzjIwMFRgYqDp37qy9PlFRUZZrjR07Vnl5eamzZ89qfRMTE23HUHDc7du3Vw899JBSSqn09HTeslbMR3fNx5MnT6pu3bopEcn/V6NGDbV9+/Yij63ImI/umY95qlatmj8fq1evrv761786fKyruOUt66ioKAkNDZV69erJoEGDJCAgQJKTk+WWW26x9Hvqqacs7ZUrV0pwcLB0795dfv755/x/bdu2lYCAANm8ebOIiGzatEmuXr0qI0eOtNwNjhkzpsixpaWlSXp6uowZM0Zuuukmy9ecubPMycmRf//739KvXz9p0KBBfh4WFiYPP/ywpKamyrlz5yzHPPHEE5ZrderUSXJycuTHH3/Mz4YOHSpKKYd++1u6dKl8/fXX8vLLLxd7/JUB87Fs56Ofn580adJEhgwZIitXrpS3335bwsLCpH///vL9998X+3uqaJiPZTsf86xbt05SUlLklVdekVtvvVUuXrxY7O+npNzylvXChQulcePG4u3tLbVq1ZImTZqIp6f1dwNvb2+pW7euJTt06JBkZmZKzZo1bc976tQpEZH8/zCNGjWyfD00NFRCQkIKHVve20PNmzd3/BsqxOnTp+XSpUvSpEkT7WtNmzaV3NxcOXbsmDRr1iw/v/XWWy398sZc8HMgR5w7d04mT54s48ePl3r16hX7+MqA+XhdWcxHEZE//OEP4u3tLR9++GF+dv/990ujRo3k2Weftbz1WBkxH68rq/mYp2vXriIi0rNnT7n//vulefPmEhAQICNGjCjReYvDLQW5ffv20q5du0L7VK1aVZuEubm5UrNmTUlKSrI9JjQ01GVjdCcvLy/bXDnxYK65c+fK1atXZeDAgfl/V3f8+HERuT6Bjxw5InXq1Cn2nzVUJMzHwrlyPv7www+yfv16WbJkiSW/+eabpWPHjrJt2zanxliRMB8L58r5eCMNGzaU1q1bS1JSUsUvyM5q2LChbNq0STp06FDoaszw8HARuf4b42/fBjl9+nSRv0U1bNhQRET27dsnUVFRN+zn6NszoaGh4ufnJ99++632tYMHD4qnp2ep3rkePXpUzpw5Y/kNM8/MmTNl5syZkpaWJq1atSq1MVRUzMfiy9vgIicnR/vatWvXJDs7u9SuXdExH10rKytLrly5UqbXNOLPnhz14IMPSk5OjkyfPl37WnZ2dv4S9aioKPHx8ZGEhATLb03z5s0r8hpt2rSRiIgImTdvnrbk/bfn8vf3F5Gil8V7eXlJjx49ZM2aNZadX06ePCnLli2Tjh07SlBQUJHjKsjRZf2jRo2S5ORky7/FixeLyPXPWZKTkyUiIqLY1wfz8bccnY+33XabeHp6yvLlyy3jP378uGzdulVat25d7GvjOubjrxydj9nZ2ba/hOzatUu+/vrrIt+pcLVydYd83333SXx8vMyaNUu+/PJL6dGjh/j4+MihQ4dk5cqVMn/+fBkwYICEhobKM888I7NmzZLevXtLTEyMpKWlybp166RGjRqFXsPT01MWLVokffr0kVatWklcXJyEhYXJwYMHZf/+/bJhwwYREWnbtq2IXC940dHR4uXlJYMGDbI954wZM2Tjxo3SsWNHefrpp8Xb21sWL14sV65ckdmzZzv1Wji6rL9NmzbSpk0bS5Y38Zs1ayb9+vVz6vpgPv6Wo/MxNDRUhg0bJn/729+kW7du0r9/fzl//ry8/vrrkpWVJZMnT3bq+mA+/paj8/HChQtSr149GThwoDRr1kz8/f3l66+/lsTERAkODpbnnnvOqes7rSyXdN9oJ5qChgwZovz9/W/49SVLlqi2bdsqX19fFRgYqFq0aKEmTJigMjIy8vvk5OSoF198UYWFhSlfX1/VpUsXtW/fPm0nmoLL+vOkpqaq7t27q8DAQOXv769atmypEhIS8r+enZ2tRo4cqUJDQ5WHh4dlib8UWNavlFJ79uxR0dHRKiAgQPn5+amuXbtqf+Zxo9fHbozOLOvPw589Xcd8dM98vHbtmkpISFCtWrVSAQEBKiAgQHXt2lV98sknRR5bkTEfy34+XrlyRY0ePVq1bNlSBQUFKR8fHxUeHq4ee+wxlZ6eXuixpcFDKRd+Eg4AAJxSrj5DBgCgoqIgAwBgAAoyAAAGoCADAGAACjIAAAagIAMAYAAKMgAABnB4py4eao/ClPWfszMfURjmI0zi6HzkDhkAAANQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADAABRkAAANQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADAABRkAAANQkAEAMAAFGQAAA1CQAQAwAAUZAAADUJABADCAt7sHAJRn3t6O/S90//33a9mdd95paa9evVrr07x5cy37/PPPtezgwYMOjQOAubhDBgDAABRkAAAMQEEGAMAAFGQAAAzgoZRSDnX08CjtsaAI9evX17LZs2drWWRkpJY1atRIy86cOeOScYmIODiNXMYd8zEwMFDLPvroIy3z9/fXsoiICC07evSopd2iRQuHxrF7924tW7p0qZZt3rxZy7799luHrlHeVYb5aIoOHTo41C8qKkrLJk6cqGWbNm2ytJOTk7U+dnP7yJEjDo3DHRydj9whAwBgAAoyAAAGoCADAGAACjIAAAZgUZehfH19tSwpKUnL7HaAevnll7VsypQprhnYDVS0RTQPPPCAltktQGnbtm2pjsNR2dnZWrZu3Tot++Mf/6hl586dK5UxuVNFm48l4ePjo2XBwcFadvnyZUt73LhxWp9BgwZp2e23365lrnz97V5bu59xkydPdtk1XY1FXQAAlCMUZAAADEBBBgDAABRkAAAMwOMXnVS1alUts9uNyW53mvfee8/SzszM1PosWbJEy+wWcH3xxRdaNmfOHC1D8XzwwQdalpub64aROObUqVNa9sknn2iZ3S5fBefosWPHXDcwuN2iRYu0LC4uTssK7hx36623ltqY8mzdulXLOnXqVOrXNRV3yAAAGICCDACAASjIAAAYgIIMAIABWNTlgAYNGmjZSy+9pGV2uzvZadeunaWdkJCg9RkwYIBD53r22We1zJWPVayszp49q2VBQUFlPxAH2T3y0c4vv/yiZdWrVy+yT2hoqJYdP35cy+x2DEPZee2117Rs2LBhWma3c1TBRVyHDh3S+tjtVmfXz24BZHx8vJa1adNGywrau3evlq1fv77I48oj7pABADAABRkAAANQkAEAMAAFGQAAA7Coq4D69etrmd0CLrtFV+fPn9eyzZs3a1nBx4SlpqZqfex2AktLS9OyjRs3ahlKzm6hnd0COjt2u2YV3AVJRGTZsmWW9lNPPaX1OXHihJa99dZbWvbVV19pmd0OcHY7NBVc1PW73/1O67NgwQIts9uZzu77RNk5efKk08cW3DVr8ODBWp+ffvrJoXO98MILWvbII49o2c0336xl3333naX9+9//XutTku/TZNwhAwBgAAoyAAAGoCADAGAACjIAAAbwUHZbtth19PAo7bGUuTp16mhZSkqKlrVo0ULLLly4oGWvv/66lj3//PNaNmLECEv7lVde0fpcvHhRy+wWRaxZs0bL3MHBaeQypT0fW7VqpWV2j7q0s2fPHoeOLbhQyu6xh3YLs+zYLUbs3r27lr3xxhtaVnChTmBgoNbHbpcykxd1VbT5WBJ2r4VdtnDhQkvb19dX69O4cWMts3tcot35v/nmGy1buXKlltktCCvvHJ2P3CEDAGAACjIAAAagIAMAYIBK9Rny7bffbmnbPTGkXr16WnbgwAEtGzNmjJZ9/PHHWmb32d7hw4ct7atXr2p9hg4dqmXLly/XMlNUtM/sPD3131XtPu/q16+fy645ZcoULbPb+KVu3bpatmTJEi2ze0KTI+zOtWvXLi2zm4+XLl1y6pquVtHmY0l07NhRy5KTk7XMbpMOR8yZM0fLPvjgAy07ePCgltmtxamI+AwZAIByhIIMAIABKMgAABiAggwAgAEq1aKugk8zuffeex06bvr06Vrm6B+vF/xjexGRJ5980tK2W0Rj9+Qfk1WGRTQtW7bUsnXr1mlZ7dq1y2I4LrN9+3ZLOyYmRutj9yQzk1WG+VgSdgta7TaSccTf/vY3LfvnP/+pZXb/r1QWLOoCAKAcoSADAGAACjIAAAagIAMAYIAKu6hr4sSJWvbSSy9Z2ufOndP6tG7dWst++OEHh645bdo0Lfvzn/+sZdu2bbO07Z6WUt5U1kU0Y8eO1bK5c+e6YSTOmz9/vqU9btw4N43EdSrrfHSU3ZOc+vbta2n3799f62O361dYWJiW5eTkaNnevXu1rODPZBGRjz76yNLOysrS+pQ3LOoCAKAcoSADAGAACjIAAAagIAMAYIAKsajLboHCjh07tKxFixaW9rVr17Q+//vf/xy6pt3rUaNGDS2ze4zf5cuXLe3OnTtrfb744guHxmGKyrqIxsfHR8tmzZqlZXaLv0wxfPhwSzs1NVXrs2/fvrIajktU1vlY2mrVqqVldjsevvXWW1oWHBzs0DVWrVplab/77rtan7Vr1zp0LlOwqAsAgHKEggwAgAEoyAAAGICCDACAASrsoq7XX39dyx588EFLu1q1ak5f0+71sHsp7RaJJSYmWtp2u9WcOXPG6bG5A4toflWlShUtW7FihZb16dOnLIZTbHZz9vDhw1r217/+Vcu+/vprLXPHgjDmo3mioqK0bNGiRVrWsGFDS9vutZ0yZYqW2S2mNAWLugAAKEcoyAAAGICCDACAASjIAAAYoEIs6nJU8+bNLW27x4bZiYyM1DK7xzvavZQFd0ESEXnjjTccum55wiKawj3zzDNa9vLLLzt1rp9++knL7BbHzJgxQ8u+/fZbLWvSpIlT47Czfv16LYuNjdWyq1evuuyadpiP5UNoaKiWPfLII5b2c889p/UJCAjQsmeffVbLXnnlFS3Lzc0tzhBdgkVdAACUIxRkAAAMQEEGAMAAFGQAAAxQqRZ1OWvDhg1a1r17dy3bu3evlrVu3bpUxmQaFtEUzm7xyttvv21px8TEOH3+5ORkLbN7VN6kSZO0rOBcvvPOO7U+ffv21bIDBw5o2R133KFldgvJtm7dqmVPPfWUpV2SxTfMx4qjQ4cOWvbpp586dGzNmjW1zNFH7LoSi7oAAChHKMgAABiAggwAgAH4DLmA3//+91pm9/lc1apVtcxuExC7DRsqIj6zK76CT4VatWqV1qcknytfunRJyzIyMrTsrbfesrTt1kIsX75cy7Kzs7UsJCSkOEO0aN++vaX9xRdfOH2u8jwfvb29tczuyXQXLlxw2TVN5uPjo2V2TxC77bbbtGzs2LFaZveUstLGZ8gAAJQjFGQAAAxAQQYAwAAUZAAADKCvHqhE7BYL2D0xpODiGxGRzz77TMvefPNN1wwMlULBJx5t27ZN61OSRV1+fn5aZrfwZdasWZa23cIvX19fLbNbfFQSU6dOtbTtNiOpDAYPHqxldk8Lmzlzppb94x//KJUxudO1a9e0LCcnx6Fj7Rbfmow7ZAAADEBBBgDAABRkAAAMQEEGAMAAlXqnLrtdXObOnevQsQMHDtSyDz74oMRjKq/K885IprBbZFi9enUtGz16tJZNmDChVMbkTl5eXk4fW57n40033aRldgv+mjZtqmVr167VstmzZ2tZamqqc4NzA7uFiLt379ay4OBgLYuLi9Oyd955xzUDKwZ26gIAoByhIAMAYAAKMgAABqAgAwBggEq9qGv79u1ads8992jZe++9p2VDhgwplTGVV+V5EU15Y7dD1owZM7Rs/PjxZTEcl9m5c6elfe+99zp9roo2HwMCArTM7pGY3bp107JffvlFy+wWtK5bt87SPnfuXHGG6DIFHzX5r3/9S+tj930W3PlOROSWW27RMrvXo7SxqAsAgHKEggwAgAEoyAAAGICCDACAASrV4xcfeughS7tFixZaH7uFAVu2bCmtIQHFlp2drWV2jw21W2gUGxurZatWrbK0H3vsMa2P3Y5hrjZt2rRSv0Z5deHCBS3r1auXlnXs2FHL7BalLlu2TMtOnDhhaT/++ONan4ILv4ojLCxMy7p3765lBXeia926tdbHbpHU4sWLtcwdC7hKgjtkAAAMQEEGAMAAFGQAAAxAQQYAwACVaqeugo/satOmjdbnrbfe0rInnnii1MZUUVS0nZEqAk9P/fdtu6zgIrE777xT69O/f38ts3uEn92iIjt2jwS8dOmSpV2SOcV8/JXdLl92/52WLFliadeuXVvrk5SUpGWnTp3SMrsFs3Y7rwUGBmpZQV999ZWWTZo0Scs2b96sZXaLdN2BnboAAChHKMgAABiAggwAgAEoyAAAGKBSLer6/vvvLW27XVxeffVVLXv//fdLbUwVBYtoYBLmY/HdfPPNlvbtt9/u0HF2C6zsdhGzY7dIbPXq1Zb2Z599pvU5efKkQ+c3BYu6AAAoRyjIAAAYgIIMAIABKtVnyCg9fGYHkzAfYRI+QwYAoByhIAMAYAAKMgAABqAgAwBgAAoyAAAGoCADAGAACjIAAAagIAMAYAAKMgAABqAgAwBgAAoyAAAGoCADAGAACjIAAAagIAMAYACHH78IAABKD3fIAAAYgIIMAIABKMgAABiAggwAgAEoyAAAGICCDACAASjIAAAYgIIMAIABKMgAABjg/wCJ6yDm5w+D/QAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 600x600 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "model.update(state.params)\n",
    "\n",
    "# plot a 3x3 grid of MNIST digits\n",
    "idxs = np.random.randint(0, len(X_test), size=(3, 3))\n",
    "fig, axes = plt.subplots(3, 3, figsize=(3 * 2, 3 * 2))\n",
    "\n",
    "for i in range(3):\n",
    "  for j in range(3):\n",
    "    logits = model(jnp.array([X_test[idxs[i, j]]]))\n",
    "    axes[i, j].imshow(X_test[idxs[i, j]], cmap=\"gray\")\n",
    "    axes[i, j].axis(\"off\")\n",
    "    axes[i, j].set_title(f\"Prediction: {jnp.argmax(logits)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome! We hope you've enjoyed this tutorial and learned the basics of NNX."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
